jax
f5288192 - Lift the restriction where lowering with only AbstractMesh required all inputs to have the same AbstractMesh. Now we only require the size of AbstractMesh to match across all inputs.

Commit

107 days ago

Lift the restriction where lowering with only AbstractMesh required all inputs to have the same AbstractMesh. Now we only require the size of AbstractMesh to match across all inputs. PiperOrigin-RevId: 810512984

References

#31381 - Remove old ROCm build code

#32043 - Lift the restriction where lowering with only AbstractMesh required all inputs to have the same AbstractMesh. Now we only require the size of AbstractMesh to match across all inputs.

#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

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

#34135 - [ROCm] update to test if there are GPU events when doing profiling on…

#587 - Mark ROCm support for gtsv2 as stable

#590 - adding abort logic to rocm/jax

#591 - Adding run_pytest_rocm.sh

Author

yashk2810

Committer

Google-ML-Automation

Parents

7952dbdf

jax f5288192 - Lift the restriction where lowering with only AbstractMesh required all inputs to have the same AbstractMesh. Now we only require the size of AbstractMesh to match across all inputs.

jax
f5288192 - Lift the restriction where lowering with only AbstractMesh required all inputs to have the same AbstractMesh. Now we only require the size of AbstractMesh to match across all inputs.