jax
e6bdd2ab - Add an environment variable to enable GPU collective cancelling.

Commit

194 days ago

Add an environment variable to enable GPU collective cancelling. When an XLA PjRt GPU client is constructed with `abort_collectives_on_failure` set to true, it aborts ongoing collectives when machines participating in the collectives fail. This CL adds an `XLA_PYTHON_CLIENT_ABORT_COLLECTIVES_ON_FAILURE` environment variable to enable the behavior. PiperOrigin-RevId: 776637541

References

#29505 - Add an environment variable to enable GPU collective cancelling.

#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