mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-24 00:36:04 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/jax/_src/pallas
History
Sergei Lebedev 5e7ad600e2 Removed the double re-exporting of Pallas GPU/TPU APIs 
jax.experimental.pallas.{gpu,tpu} now import directly from the relevant
jax._src.pallas.{triton,mosaic} submodules.

PiperOrigin-RevId: 641875127
2024-06-10 05:59:09 -07:00
..
mosaic
Removed the double re-exporting of Pallas GPU/TPU APIs
2024-06-10 05:59:09 -07:00
mosaic_gpu
Updated remaining `mgpu.once usages to use mgpu.single_thread`
2024-06-05 08:31:36 -07:00
triton
Removed the double re-exporting of Pallas GPU/TPU APIs
2024-06-10 05:59:09 -07:00
__init__.py
[Pallas] Upstream pallas to JAX
2023-08-01 16:43:13 -07:00
BUILD
Removed the double re-exporting of Pallas GPU/TPU APIs
2024-06-10 05:59:09 -07:00
core.py
Add dynamic grid support to emit_pipeline
2024-04-09 23:55:51 -07:00
pallas_call.py
Merge pull request #21298 from oliverdutton:pallas_interpreter_indexing_fix
2024-06-07 12:29:31 -07:00
primitives.py
[pallas] align interpreter load/store for masked OOB slicing
2024-06-04 23:19:55 +01:00
utils.py
Reverts c2a3c0bb80434d89053b43648f88eba22b9bf1fa
2024-06-05 07:50:58 -07:00