Adam Paszke f1ab7514db Make sure we take libTPU version into account in the Pallas lowering ...

Also, strengthen the presubmit to make sure we catch more errors.

PiperOrigin-RevId: 726061633

2025-02-12 08:15:57 -08:00

..

mosaic

Make sure we take libTPU version into account in the Pallas lowering

2025-02-12 08:15:57 -08:00

mosaic_gpu

[Pallas/Mosaic GPU] Enable progressive lowering for integer addition.

2025-02-12 06:29:25 -08:00

triton

[pallas:triton] Lift dot_general restriction on minimal tile size for a.

2025-02-11 06:27:16 -08:00

__init__.py

[Pallas] Upstream pallas to JAX

2023-08-01 16:43:13 -07:00

BUILD

[Pallas] Refactor Pallas HLO interpret mode to a standalone file.

2025-01-27 17:52:27 -08:00

core.py

Merge pull request #26423 from gnecula:debug_info_jaxpr_7

2025-02-10 12:58:26 -08:00

cost_estimate.py

[better_errors] Continue adding debug info to Jaxprs (step 7)

2025-02-09 18:14:33 +02:00

helpers.py

[Pallas TPU] Add helpers for writing collectives

2025-02-04 15:39:10 -08:00

hlo_interpreter.py

[better_errors] Continue adding debug info to Jaxprs (step 3)

2025-02-06 16:26:49 +02:00

pallas_call.py

[Pallas/Mosaic GPU] Add initial support for warpgroup semantics in lowering.

2025-02-12 01:47:49 -08:00

primitives.py

[better_errors] Continue adding debug info to Jaxprs (step 7)

2025-02-09 18:14:33 +02:00

utils.py

[Pallas TPU] Add helpers for writing collectives

2025-02-04 15:39:10 -08:00