Sergei Lebedev d34c40f6b6 [mosaic_gpu] Added a serialization pass ...

The pass adds versioning to the Mosaic GPU IR in the lowered custom calls
and can apply forward/backward migration rules. Currently, no rules are
necessary since we are at version 1.

PiperOrigin-RevId: 716596848

2025-01-17 03:12:51 -08:00

..

mosaic

[Pallas TPU] Add helpers file with copy_ref function

2025-01-15 18:34:58 -08:00

mosaic_gpu

[mosaic_gpu] Added a serialization pass

2025-01-17 03:12:51 -08:00

triton

Rename out_type -> out_sharding parameter on einsum

2025-01-16 18:16:52 -08:00

__init__.py

[Pallas] Upstream pallas to JAX

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

BUILD

[Pallas] Add empty/empty_like helper functions

2025-01-08 10:49:11 -08:00

core.py

Merge pull request #25876 from gnecula:debug_info_3

2025-01-15 09:35:03 -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] Add empty/empty_like helper functions

2025-01-08 10:49:11 -08:00

pallas_call.py

[better_errors] More cleanup

2025-01-15 10:22:29 +00:00

primitives.py

[Pallas TPU] Add vector support to pl.debug_print

2025-01-13 13:22:21 -08:00

utils.py

[Pallas TPU] Add lowering for lax.nextafter

2024-10-31 17:34:38 -07:00