Ayaka 9ba1fd2801 [Pallas TPU] Add vector support to pl.debug_print ...

PiperOrigin-RevId: 715085454

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

..

mosaic

[Pallas TPU] Add vector support to pl.debug_print

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

mosaic_gpu

[Mosaic GPU] Fix layout API bugs.

2025-01-13 12:59:30 -08:00

triton

Reverts dbe9ccd6dccd83c365021677c7e17e843d4559c4

2025-01-10 03:40:18 -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

[Pallas] Improvements to core_map

2025-01-07 16:18:30 -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

Cleanup: remove uses of no-op raise_to_shaped

2024-12-12 09:49:06 -08: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