jax
66fc5e88 - [Mosaic GPU] Add support for arbitrary reductions of tiled layouts

Commit

262 days ago

[Mosaic GPU] Add support for arbitrary reductions of tiled layouts This significantly generalizes our ability to perform reductions, to the point where pretty much all tiled layouts can be handled out of the box. The code is slightly longer than the few special cases we've implemented in the past, but overall is much more general. This also includes a hypothesis test that verifies that we always return the right answers, even for randomly sampled layouts. The test could still be improved in that we'll skip all the cases where we fail to synthesize a load/store for the layout, but it's already caught a number of problems in the initial implementation. PiperOrigin-RevId: 755303138

References

#28448 - [Mosaic GPU] Add support for arbitrary reductions of tiled layouts

#477 - CI: 06/20/25 upstream sync

#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 - [ROCm] Resolve undefined behavior in bitshift unit test

#579 - Create rocm-test-requirements.txt

#581 - Fix/pallas tests shared memory

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

#585 - update a test for checking zero ROCm GPU event