Sharad Vikram 8dfbf90602 [Pallas/Mosaic] Add support for barrier semaphores ...

PiperOrigin-RevId: 586289340

2023-11-29 04:04:11 -08:00

..

mosaic

[Pallas/Mosaic] Add support for barrier semaphores

2023-11-29 04:04:11 -08:00

triton

Increase minimum jaxlib version to 0.4.19.

2023-11-17 09:38:31 -08:00

__init__.py

[Pallas] Upstream pallas to JAX

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

BUILD

Load Python rules from rules_python

2023-08-24 10:22:57 -07:00

core.py

[Pallas TPU] Add support for hoisted scratch spaces

2023-10-24 17:29:54 -07:00

indexing.py

Switch from flake8 to Ruff

2023-11-15 22:35:52 -05:00

pallas_call.py

[Pallas] Allow interpret mode on non-CPU backends if backend-specific lowerings are not registered

2023-11-16 12:46:43 -08:00

primitives.py

[pallas] dot fixes.

2023-10-16 12:35:43 -07:00

utils.py

[Pallas] Add support for lax.fori_loop to Mosaic lowering

2023-08-04 13:43:43 -07:00