jax
142ace2c - Move jax._src.callback to its own BUILD rule

Commit

220 days ago

Move jax._src.callback to its own BUILD rule Creating smaller build rules enforces better organized dependency graphs in the JAX project, helps pytype propagate annotations correctly, prevents use of internal APIs, and leads to improved build and iteration times. This will unblock moving `jax/_src/lax` to its own BUILD rule. PiperOrigin-RevId: 770980397

References

#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

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