jax
123022ca - [Mosaic] Make `tpu.relayout` an explicit operation, merge in existing behavior, stop calling relayout() in apply.

Commit

253 days ago

[Mosaic] Make `tpu.relayout` an explicit operation, merge in existing behavior, stop calling relayout() in apply. This change should reduce complexity and make it easier to see what happened in a graph. Note - there are still cases where certain relayout() calls are not ops yet, those will be migrated in the future. Specifically, see the note in the CL around force_relayout. Added helper methods to generate full like vectors. Followup for subsequent CLs: Simplify Relayout rule in future CLs, maybe break up into smaller sub relayouts with nice names. Followup for subsequent CLs: Unify transpose in here PiperOrigin-RevId: 758240646

References

#28663 - [Mosaic] Make `tpu.relayout` an explicit operation, merge in existing behavior, stop calling relayout() in apply.

#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 - 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