jax
b09724a9 - [Mosaic] Allow vector::Extract for non-32 bits vector result.

Commit

163 days ago

[Mosaic] Allow vector::Extract for non-32 bits vector result. It should have been supported in apply. This allows us to run batched non-32 bits matmuls. https://github.com/jax-ml/jax/blob/0a572631b8791b08a975c9a7087f801a02f5d1ee/jaxlib/mosaic/dialect/tpu/transforms/canonicalize_mosaic.cc#L504-L509 PiperOrigin-RevId: 786746709

References

#30371 - [Mosaic] Allow vector::Extract for non-32 bits vector result.

#31381 - Remove old ROCm build code

#31720 - Fix ann_test.py numerical bug in target reshape

#31768 - [ROCm] Support lowering through PJRT_Triton_Extension

#32115 - Relax version requirements for ROCm Jax Plugin wheels

#33157 - Resolve undefined behavior in bitshift unit test

#33186 - Make nvidia version data optional for ROCm builds

#579 - Create rocm-test-requirements.txt

#580 - GESVDJ support for ROCm GPUs in JAX

#581 - Fix/pallas tests shared memory

#584 - Use plain bazel to test jax, use hermetic rocm dependency