rocm_jax

mirror of https://github.com/ROCm/jax.git synced 2025-04-22 22:06:05 +00:00

History

Fix invalid lowerings for ROCm in Pallas (#223 )

popcount and clz were effectively broken on ROCm,
since math_dialect had incorrect lowerings.

Use the device intrinsics for these functions, as
well as for exp and absf, which fixes some accuracy issues in
the pallas tests.

Docs for OCML/OCKL

- https://github.com/ROCm/llvm-project/blob/amd-staging/amd/device-libs/doc/OCML.md
- https://github.com/ROCm/llvm-project/blob/amd-staging/amd/device-libs/doc/OCKL.md

2025-02-14 11:27:52 -06:00

mosaic

Reverts f1b894d14a28ac22a037fb79177b991275c75a18

2025-01-17 07:00:31 -08:00

mosaic_gpu

[mosaic_gpu] Added a serialization pass

2025-01-17 03:12:51 -08:00

triton

Fix invalid lowerings for ROCm in Pallas (#223 )

2025-02-14 11:27:52 -06:00

__init__.py

[Pallas] Upstream pallas to JAX