mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-18 04:46:06 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/jax/_src/pallas/triton
History
Sergei Lebedev 498e81ab10 Pallas now exclusively uses XLA for compiling kernels on GPU 
The old lowering pass via Triton Python APIs has been removed and the
JAX_TRITON_COMPILE_VIA_XLA environment variable no longer has any effect.

PiperOrigin-RevId: 621857046
2024-04-04 07:47:26 -07:00
..
__init__.py
Register GPU/TPU lowering for pallas_call_p lazily
2024-04-01 14:40:33 -07:00
BUILD
Pallas now exclusively uses XLA for compiling kernels on GPU
2024-04-04 07:47:26 -07:00
lowering.py
Allows for splitting the transpose of a scan into a scan and a map.
2024-03-28 10:54:50 -07:00
pallas_call_registration.py
Pallas now exclusively uses XLA for compiling kernels on GPU
2024-04-04 07:47:26 -07:00