jax
f4626f42 - Update ManualComputation round-tripping to not have round-tripping attrs on the CallOps.

Commit

270 days ago

Update ManualComputation round-tripping to not have round-tripping attrs on the CallOps. This is needed as sometimes the op has `stablehlo.token` types, and during MLIR<->HLO round-tripping, MLIR type conversion converts this to `mhlo.token` and "accidentally" discards the unregistered attrs (the `frontend_attributes`) - note MLIR has no guarantees about preserving these, so we can't submit a fix to MLIR. This happens because StableHLO->HLO conversion still does StableHLO->MHLO conversion, but this should be removed soon. To unblock ourselves instead of waiting for StableHLO to remove the intermediate MHLO pass, we can just move the frontend attrs to the GlobalToLocal/LocalToGlobal custom calls. PiperOrigin-RevId: 753995222

References

#28000 - Update ManualComputation round-tripping to not have round-tripping attrs on the CallOps.

#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