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

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: 897706398

Author: danielsuo

chex/_src/asserts_chexify_test.py

@@ -37,6 +37,18 @@
 chexify_sync = functools.partial(asserts_chexify.chexify, async_check=False)
 
 
+def _device_put_replicated(x, devices):
+  mesh = jax.sharding.Mesh(
+      np.array(devices), axis_names=('_device_put_replicated',)
+  )
+  sharding = jax.sharding.NamedSharding(
+      mesh, jax.sharding.PartitionSpec('_device_put_replicated')
+  )
+  return jax.tree_util.tree_map(
+      lambda v: jax.device_put(np.stack([v] * len(devices)), sharding), x
+  )
+
+
 def get_chexify_err_regex(name, msg):
   return re.escape(_ai.get_chexify_err_message(name, 'ANY')).replace(
       'ANY', f'.*{msg}.*'
@@ -401,8 +413,7 @@ def run_test_suite_with_log_abs_fn(self, make_log_fn, jax_transform, devices,
             'c': np.array([[5, -1] for _ in range(10)])
         }
     }
-    (x_pos, x_with_neg) = jax.device_put_replicated((x_pos, x_with_neg),
-                                                    devices)
+    x_pos, x_with_neg = _device_put_replicated((x_pos, x_with_neg), devices)
 
     all_valid_args = ((x_pos, x_pos),)
     all_invalid_args = (

chex/_src/asserts_test.py

@@ -44,6 +44,18 @@ def emplace(arrays, dtype):
   return jnp.array(arrays, dtype=dtype)
 
 
+def _device_put_replicated(x, devices):
+  mesh = jax.sharding.Mesh(
+      np.array(devices), axis_names=('_device_put_replicated',)
+  )
+  sharding = jax.sharding.NamedSharding(
+      mesh, jax.sharding.PartitionSpec('_device_put_replicated')
+  )
+  return jax.tree_util.tree_map(
+      lambda v: jax.device_put(np.stack([v] * len(devices)), sharding), x
+  )
+
+
 class AssertsSwitchTest(parameterized.TestCase):
   """Tests for enable/disable_asserts."""
 
@@ -1214,7 +1226,7 @@ def test_assert_tree_is_on_host(self):
 
     # Check sharded Jax arrays on CPUs.
     asserts.assert_tree_is_on_host(
-        {'a': jax.device_put_replicated(np.zeros(1), (cpu,))},
+        {'a': _device_put_replicated(np.zeros(1), (cpu,))},
         allow_cpu_device=True,
         allow_sharded_arrays=True,
     )
@@ -1242,7 +1254,7 @@ def test_assert_tree_is_on_host(self):
         AssertionError, _get_err_regex("'a' resides on.*CPU.*disallowed")
     ):
       asserts.assert_tree_is_on_host(
-          {'a': jax.device_put_replicated(np.zeros(1), (cpu,))},
+          {'a': _device_put_replicated(np.zeros(1), (cpu,))},
           allow_cpu_device=False,
       )
 
@@ -1251,7 +1263,7 @@ def test_assert_tree_is_on_host(self):
         AssertionError, _get_err_regex("'a' resides on.*CPU.*disallowed")
     ):
       asserts.assert_tree_is_on_host(
-          {'a': jax.device_put_replicated(np.zeros(1), (cpu,))},
+          {'a': _device_put_replicated(np.zeros(1), (cpu,))},
           allow_cpu_device=False,
           allow_sharded_arrays=True,
       )
@@ -1352,7 +1364,7 @@ def _format(*devs):
     # a "sharded" array (device_set has length 1). The array is treated as a
     # regular single-device array.
     cpu = jax.local_devices(backend='cpu')[0]
-    cpu_tree = jax.device_put_replicated(np_tree, (cpu,))
+    cpu_tree = _device_put_replicated(np_tree, (cpu,))
 
     # Single-device array is NOT considered sharded.
     with self.assertRaisesRegex(
@@ -1366,10 +1378,10 @@ def _format(*devs):
     if _num_devices_available('tpu') > 1:
       tpu_1, tpu_2 = jax.devices('tpu')[:2]
 
-      tpu_1_tree = jax.device_put_replicated(np_tree, (tpu_1,))
-      tpu_2_tree = jax.device_put_replicated(np_tree, (tpu_2,))
-      tpu_1_2_tree = jax.device_put_replicated(np_tree, (tpu_1, tpu_2))
-      tpu_2_1_tree = jax.device_put_replicated(np_tree, (tpu_2, tpu_1))
+      tpu_1_tree = _device_put_replicated(np_tree, (tpu_1,))
+      tpu_2_tree = _device_put_replicated(np_tree, (tpu_2,))
+      tpu_1_2_tree = _device_put_replicated(np_tree, (tpu_1, tpu_2))
+      tpu_2_1_tree = _device_put_replicated(np_tree, (tpu_2, tpu_1))
 
       asserts.assert_tree_is_sharded(tpu_1_2_tree, devices=(tpu_1, tpu_2))
       asserts.assert_tree_is_sharded(tpu_2_1_tree, devices=(tpu_2, tpu_1))

chex/_src/variants.py

@@ -28,6 +28,7 @@
 import jax
 from jax import tree_util
 import jax.numpy as jnp
+import numpy as np
 import toolz
 
 FLAGS = flags.FLAGS
@@ -550,7 +551,9 @@ def bcast_fn(x):
       x = jnp.asarray(x)
       x = jnp.broadcast_to(x, (n_devices_,) + x.shape)
       if not isinstance(x, jax.core.Tracer):
-        return jax.device_put_sharded(list(x), devices_)
+        mesh = jax.sharding.Mesh(np.array(devices_), ("_device_put_sharded",))
+        sharding = jax.NamedSharding(mesh, jax.P("_device_put_sharded"))
+        return jax.device_put(jnp.stack(list(x)), sharding)
       return x
 
     if broadcast_args_to_devices: