jax authors e9ce8fb92d Merge pull request #27227 from jburnim:jburnim_pallas_interpret_mode4 ...

PiperOrigin-RevId: 738235363

2025-03-18 20:22:27 -07:00

..

__init__.py

Removed the double re-exporting of Pallas GPU/TPU APIs

2024-06-10 05:59:09 -07:00

BUILD

[Pallas] Enable skipping of floating-point operations when interpreting Pallas TPU kernels on CPU.

2025-03-10 15:14:00 -07:00

core.py

[pallas] pallas_call_p is now parameterized by a mesh

2025-03-17 16:30:40 -07:00

error_handling.py

Propagate source ranges in location.

2025-02-03 17:32:59 -08:00

helpers.py

[Pallas TPU] Add helpers for writing collectives

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

interpret.py

Adds option to initialize buffers to NaNs or zeros in TPU interpret mode.

2025-03-18 12:24:45 -07:00

lowering.py

Better mosaic lowering for dynamic shapes, extend an interpreter into shape_poly dimexpr and lower them alongside the graph if we are in a dynamic export regime.

2025-03-18 15:51:15 -07:00

pallas_call_registration.py

[pallas] pallas_call_p is now parameterized by a mesh

2025-03-17 16:30:40 -07:00

pipeline.py

[Pallas TPU] Hoist prologue and epilogue outside of pipeline loop

2025-03-18 09:40:43 -07:00

primitives.py

[pallas] Support DMA start partial discharge and run_scoped() does its own partial discharge.

2025-02-01 08:23:23 -08:00

random.py

[Pallas] Add legacy PRNG key support to Pallas PRNG

2025-03-14 12:30:04 -07:00

verification.py

[Mosaic TPU] Add support for modeling loads/stores and fix minor issues in model extraction

2024-12-05 07:06:19 -08:00