jax

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

[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