mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-24 15:36:07 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/jax/_src/pallas
History
Ayaka 1dff3a2c71 [Pallas GPU] Add 32-bit lowering rule for lax.erf_inv 
Add 32-bit lowering rule for `lax.erf_inv` for Pallas GPU, and move the original TPU test case into the general test

PiperOrigin-RevId: 668681910
2024-08-28 17:59:15 -07:00
..
mosaic
[Pallas GPU] Add 32-bit lowering rule for lax.erf_inv
2024-08-28 17:59:15 -07:00
mosaic_gpu
Clean up BUILD files.
2024-08-26 09:11:17 -07:00
triton
[Pallas GPU] Add 32-bit lowering rule for lax.erf_inv
2024-08-28 17:59:15 -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