flax

355e141a - [pmap] In-line definitions of `jax.device_put_sharded` and `jax.device_put_replicated`.

[pmap] In-line definitions of `jax.device_put_sharded` and `jax.device_put_replicated`. NOTE: Added a `axis_name` argument because `jax_utils` is a commonly used utility and it may be beneficial for `flax` users to be able to specify the axis name to match the axis name used in their experiment. Both `jax.device_put_sharded` and `jax.device_put_replicated` were deprecated in JAX v0.8.1 in November 2025. We in-line their definitions using public JAX APIs taking the `jax_pmap_shmap_merge=True` branch, which was made the default in JAX v0.8.0 in October 2025. Please see the below for more information: - JAX CHANGELOG: https://docs.jax.dev/en/latest/changelog.html - Migrating from `jax.pmap`: https://docs.jax.dev/en/latest/migrate_pmap.html PiperOrigin-RevId: 895857088