jax
5ff55080 - [mutable-arrays] add discharge_state2 for ClosedJaxprs, with caching

Commit

127 days ago

[mutable-arrays] add discharge_state2 for ClosedJaxprs, with caching This helper function follows our usual pattern of caching on ClosedJaxprs, which handles jaxprs and const lists together by id, and also keeps weak references to both. We can probably update most or all callers to use discharge_state2 instead of discharge_state, at which point we can delete the latter. This PR only updates the jit lowering rule, jit discharge rule, and scan discharge rule to use the new helper.

References

#31131 - [mutable-arrays] add discharge_state2 for ClosedJaxprs, with caching

#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