jax
5b1ddeff - Don't serialize MLIR modules unnecessarily when passing them to the Python bindings.

Commit

101 days ago

Don't serialize MLIR modules unnecessarily when passing them to the Python bindings. We used to do this because the MLIR generated by the Python MLIR extensions belonged to a different C++ extension than the compiler. But since we switched over to using the pywrap rules this is no longer true and we can just pass the unserialized module. PiperOrigin-RevId: 811346455

References

#31381 - Remove old ROCm build code

#32056 - Don't serialize MLIR modules unnecessarily when passing them to the Python bindings.

#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…