jax
ef1b3e92 - [JAX] Make fully replicated sharding to avoid materializing the same host buffers

Commit

234 days ago

[JAX] Make fully replicated sharding to avoid materializing the same host buffers This change recognizes fully replicated shardings (based on JAX sharding), and materialize only one host buffer. This saves the cost of repeatedly materializing the same host buffer for multiple devices, and makes it streamlined to create a multi-device IFRT array. Clean up `JAX_IFRT_VERSION_NUMBER < 2` since we are well past it. PiperOrigin-RevId: 758320113

References

#28493 - [JAX] Make fully replicated sharding to avoid materializing the same host buffers

#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

#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