jax authors 6ad38af473 Merge pull request #26368 from ROCm:fix-rocm-pallas-lowerings ...

PiperOrigin-RevId: 724328946

2025-02-07 06:52:01 -08:00

..

mosaic

Start a new TPU interpret mode for Pallas.

2025-02-06 13:04:14 -08:00

mosaic_gpu

[pallas:mgpu] Cast all indices to i32 during lowering.

2025-02-05 07:37:04 -08:00

triton

Fix invalid lowerings for ROCm in Pallas

2025-02-06 13:59:58 -06: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 #26078 from gnecula:debug_info_jaxpr

2025-02-04 10:54:26 -08:00

cost_estimate.py

[Pallas] Add support for run_state to cost estimator.

2024-12-06 10:36:02 -08: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

Start a new TPU interpret mode for Pallas.

2025-02-06 13:04:14 -08:00

primitives.py

[pallas/pallas_mgpu] Discharging run_scoped should not be discharging the intermediates

2025-02-01 09:37:03 -08:00

utils.py

[Pallas TPU] Add helpers for writing collectives

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