Chris Jones 7b161fb76c [jax:pallas] Use 64-bit indexing when necessary when lowering to Triton. ...

PiperOrigin-RevId: 670530776

2024-09-03 06:20:39 -07:00

..

mosaic

Disable scoped tracing in pipelining

2024-08-30 01:28:19 -07:00

mosaic_gpu

Clean up BUILD files.

2024-08-26 09:11:17 -07:00

triton

[jax:pallas] Use 64-bit indexing when necessary when lowering to Triton.

2024-09-03 06:20:39 -07:00

__init__.py

[Pallas] Upstream pallas to JAX

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

BUILD

Clean up BUILD files.

2024-08-26 09:11:17 -07:00

core.py

[Pallas] Add interpret mode rules for semaphores (local signal, wait, read, DMAs).

2024-08-21 11:11:11 -07:00

pallas_call.py

[Pallas] Add explicit TPU compiler params interface with docstrings.

2024-08-20 15:38:38 -07:00

primitives.py

Support strided load / store in interpret mode

2024-08-28 05:10:54 -07:00

utils.py

[Pallas GPU] Add 32-bit lowering rule for lax.erf_inv

2024-08-28 17:59:15 -07:00