[pmap] Deprecate jax.sharding.PmapSharding, jax.device_put_replicated and jax.device_put_sharded.

PiperOrigin-RevId: 831569381

Author: danielsuo

CHANGELOG.md

@@ -49,6 +49,12 @@ When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.
     GPU (({jax-issue}`#33062`).
 
 * Deprecations:
+  * `jax.sharding.PmapSharding` is now deprecated. Please use
+    `jax.NamedSharding` instead.
+  * `jx.device_put_replicated` is now deprecated. Please use `jax.device_put`
+    with the appropriate sharding instead.
+  * `jax.device_put_sharded` is now deprecated. Please use `jax.device_put` with
+    the appropriate sharding instead.
   * Default `axis_types` of `jax.make_mesh` will change in JAX v0.9.0 to return
   `jax.sharding.AxisType.Explicit`. Leaving axis_types unspecified will raise a
   `DeprecationWarning`.

jax/__init__.py

@@ -97,8 +97,8 @@
 from jax._src.xla_bridge import device_count as device_count
 from jax._src.api import device_get as device_get
 from jax._src.api import device_put as device_put
-from jax._src.api import device_put_sharded as device_put_sharded
-from jax._src.api import device_put_replicated as device_put_replicated
+from jax._src.api import device_put_sharded as _deprecated_device_put_sharded
+from jax._src.api import device_put_replicated as _deprecated_device_put_replicated
 from jax._src.xla_bridge import devices as devices
 from jax._src.api import disable_jit as disable_jit
 from jax._src.api import eval_shape as eval_shape
@@ -197,6 +197,16 @@
     "jax.ArrayRef is deprecated; use jax.Ref instead.",
     Ref
   ),
+  # Added for v0.8.1
+  "device_put_replicated": (
+    "jax.device_put_replicated is deprecated; use jax.device_put instead.",
+    _deprecated_device_put_replicated
+  ),
+  # Added for v0.8.1
+  "device_put_sharded": (
+    "jax.device_put_sharded is deprecated; use jax.device_put instead.",
+    _deprecated_device_put_sharded
+  ),
   # Finalized 2025-03-25; remove after 2025-06-25
   "treedef_is_leaf": (
     "jax.treedef_is_leaf was removed in JAX v0.6.0: use jax.tree_util.treedef_is_leaf.",
@@ -238,6 +248,8 @@
 if _typing.TYPE_CHECKING:
   array_ref = new_ref
   ArrayRef = Ref
+  device_put_replicated = _deprecated_device_put_replicated
+  device_put_sharded = _deprecated_device_put_sharded
 else:
   from jax._src.deprecations import deprecation_getattr as _deprecation_getattr
   __getattr__ = _deprecation_getattr(__name__, _deprecations)
@@ -247,3 +259,5 @@
 import jax.lib  # TODO(phawkins): remove this export.  # noqa: F401
 
 # trailer
+del _deprecated_device_put_sharded
+del _deprecated_device_put_replicated

jax/_src/api.py

@@ -2871,23 +2871,23 @@ def device_put_sharded(shards: Sequence[Any], devices: Sequence[xc.Device]):  #
     >>> import jax
     >>> devices = jax.local_devices()
     >>> x = [jax.numpy.ones(5) for device in devices]
-    >>> y = jax.device_put_sharded(x, devices)
-    >>> np.allclose(y, jax.numpy.stack(x))
+    >>> y = jax.device_put_sharded(x, devices)  # doctest: +SKIP
+    >>> np.allclose(y, jax.numpy.stack(x))  # doctest: +SKIP
     True
 
     Passing a list of nested container objects with arrays at the leaves for
     ``shards`` corresponds to stacking the shards at each leaf. This requires
     all entries in the list to have the same tree structure:
 
     >>> x = [(i, jax.numpy.arange(i, i + 4)) for i in range(len(devices))]
-    >>> y = jax.device_put_sharded(x, devices)
-    >>> type(y)
+    >>> y = jax.device_put_sharded(x, devices)  # doctest: +SKIP
+    >>> type(y)  # doctest: +SKIP
     <class 'tuple'>
-    >>> y0 = jax.device_put_sharded([a for a, b in x], devices)
-    >>> y1 = jax.device_put_sharded([b for a, b in x], devices)
-    >>> np.allclose(y[0], y0)
+    >>> y0 = jax.device_put_sharded([a for a, b in x], devices)  # doctest: +SKIP
+    >>> y1 = jax.device_put_sharded([b for a, b in x], devices)  # doctest: +SKIP
+    >>> np.allclose(y[0], y0)  # doctest: +SKIP
     True
-    >>> np.allclose(y[1], y1)
+    >>> np.allclose(y[1], y1)  # doctest: +SKIP
     True
 
   See Also:
@@ -2953,8 +2953,8 @@ def device_put_replicated(x: Any, devices: Sequence[xc.Device]):  # noqa: F811
     >>> import jax
     >>> devices = jax.local_devices()
     >>> x = jax.numpy.array([1., 2., 3.])
-    >>> y = jax.device_put_replicated(x, devices)
-    >>> np.allclose(y, jax.numpy.stack([x for _ in devices]))
+    >>> y = jax.device_put_replicated(x, devices)  # doctest: +SKIP
+    >>> np.allclose(y, jax.numpy.stack([x for _ in devices]))  # doctest: +SKIP
     True
 
   See Also:

jax/sharding.py

@@ -19,7 +19,7 @@
 from jax._src.sharding_impls import (
     NamedSharding as NamedSharding,
     SingleDeviceSharding as SingleDeviceSharding,
-    PmapSharding as PmapSharding,
+    PmapSharding as _deprecated_PmapSharding,
     set_mesh as set_mesh,
 )
 from jax._src.partition_spec import (
@@ -39,3 +39,21 @@
     auto_axes as auto_axes,
     explicit_axes as explicit_axes,
 )
+
+_deprecations = {
+  # Added for v0.8.1
+  "PmapSharding": (
+    "jax.sharding.PmapSharding is deprecated; use jax.sharding.NamedSharding instead.",
+    _deprecated_PmapSharding
+  ),
+}
+
+import typing as _typing
+if _typing.TYPE_CHECKING:
+  PmapSharding = _deprecated_PmapSharding
+else:
+  from jax._src.deprecations import deprecation_getattr as _deprecation_getattr
+  __getattr__ = _deprecation_getattr(__name__, _deprecations)
+  del _deprecation_getattr
+del _typing
+del _deprecated_PmapSharding

tests/array_test.py

@@ -946,6 +946,7 @@ def test_uneven_shard_error(self):
         r"factors: \[4, 2\] should evenly divide the shape\)"):
       mps.shard_shape((8, 3))
 
+  @jtu.ignore_warning(category=DeprecationWarning)
   @jtu.thread_unsafe_test()  # cache_info isn't thread-safe
   def test_pmap_sharding_hash_eq(self):
     if jax.device_count() < 2:
@@ -1056,6 +1057,7 @@ def test_pspec_tuple(self):
       ('sharded_dim_2', (4, 2, 4), 2),
       ('sharded_dim_1_1', (2, 4), 1)
   )
+  @jtu.ignore_warning(category=DeprecationWarning)
   def test_default_pmap_sharding(self, shape, sharded_dim):
     if jax.device_count() < 4:
       self.skipTest('Test needs >= 4 devices.')
@@ -1081,6 +1083,7 @@ def test_default_pmap_sharding(self, shape, sharded_dim):
       self.assertEqual(actual_sharding.sharding_spec, expected_sharding.sharding_spec)
       self.assertEqual(actual_sharding._device_assignment, expected_sharding._device_assignment)
 
+  @jtu.ignore_warning(category=DeprecationWarning)
   def test_default_pmap_sharding_with_devices(self):
     if jax.device_count() < 4:
       self.skipTest('Test needs >= 4 devices.')
@@ -1090,6 +1093,7 @@ def test_default_pmap_sharding_with_devices(self):
     ps = jax.sharding.PmapSharding.default((4, 2), devices=new_order)
     self.assertEqual(ps._device_assignment, new_order)
 
+  @jtu.ignore_warning(category=DeprecationWarning)
   def test_default_pmap_sharding_replicated(self):
     x = np.zeros((len(jax.local_devices()), 8), dtype=np.float32)
     x = jax.pmap(lambda x: x, in_axes=0, out_axes=None, axis_name='x')(x)

tests/debugging_primitives_test.py

@@ -1105,6 +1105,7 @@ def test_visualize_wide_array(self):
     """)
     self.assertEqual(output(), expected)
 
+  @jtu.ignore_warning(category=DeprecationWarning)
   def test_visualize_pmap_sharding(self):
     ss = pxla.ShardingSpec(
         sharding=(pxla.Unstacked(8),),

tests/multiprocess/all_gather_test.py

@@ -30,6 +30,7 @@ class AllGatherTest(jt_multiprocess.MultiProcessTest):
   @parameterized.parameters(
       (np.int32,), (jnp.float32,), (jnp.float16,), (jnp.bfloat16,)
   )
+  @jtu.ignore_warning(category=DeprecationWarning)
   def test_all_gather(self, dtype):
     f = jax.pmap(lambda x: lax.all_gather(x, "i"), axis_name="i")
     xs = randint_sample(

tests/multiprocess/all_reduce_test.py

@@ -27,6 +27,7 @@ def randint_sample(shape):
 
 class AllReduceTest(jt_multiprocess.MultiProcessTest):
 
+  @jtu.ignore_warning(category=DeprecationWarning)
   def test_psum_simple(self):
     f = jax.pmap(lambda x: lax.psum(x, "i"), "i", devices=jax.devices())
     np.testing.assert_array_equal(
@@ -37,6 +38,7 @@ def test_psum_simple(self):
   @parameterized.parameters(
       (np.int32,), (jnp.float32,), (jnp.float16,), (jnp.bfloat16,)
   )
+  @jtu.ignore_warning(category=DeprecationWarning)
   def test_psum(self, dtype):
     f = jax.pmap(lambda x: lax.psum(x, "i"), axis_name="i")
     xs = randint_sample(
@@ -47,6 +49,7 @@ def test_psum(self, dtype):
     for actual in out:
       jtu.check_close(actual, expected)
 
+  @jtu.ignore_warning(category=DeprecationWarning)
   def test_psum_subset_devices(self):
     f = jax.pmap(
         lambda x: lax.psum(x, "i"), axis_name="i", devices=jax.local_devices()
@@ -57,6 +60,7 @@ def test_psum_subset_devices(self):
     for actual in out:
       np.testing.assert_array_equal(actual, expected)
 
+  @jtu.ignore_warning(category=DeprecationWarning)
   def test_psum_del(self):  # b/171945402
     f = jax.pmap(lambda x: lax.psum(x, "i"), axis_name="i")
     g = jax.pmap(lambda x: lax.psum(x, "i"), axis_name="i")
@@ -73,6 +77,7 @@ def test_psum_del(self):  # b/171945402
     for actual in out:
       np.testing.assert_array_equal(actual, expected)
 
+  @jtu.ignore_warning(category=DeprecationWarning)
   def test_psum_multiple_operands(self):
     f = jax.pmap(lambda x: lax.psum(x, "i"), axis_name="i")
     xs = randint_sample([jax.process_count(), jax.local_device_count(), 100])
@@ -85,6 +90,7 @@ def test_psum_multiple_operands(self):
     for actual in out_ys:
       np.testing.assert_array_equal(actual, expected_ys)
 
+  @jtu.ignore_warning(category=DeprecationWarning)
   def test_psum_axis_index_groups(self):
     devices = list(range(jax.device_count()))
     axis_index_groups = [devices[0::2], devices[1::2]]

tests/multiprocess/all_to_all_test.py

@@ -30,6 +30,7 @@ class AllToAllTest(jt_multiprocess.MultiProcessTest):
 
   @parameterized.parameters((np.int32,), (jnp.float32,), (jnp.float16,),
                             (jnp.bfloat16,))
+  @jtu.ignore_warning(category=DeprecationWarning)
   def test_all_to_all(self, dtype):
     f = jax.pmap(
         lambda x: lax.all_to_all(x, "i", split_axis=0, concat_axis=0),

tests/multiprocess/multihost_utils_test.py

@@ -419,6 +419,7 @@ def test_host_local_to_global_reshard_committed_single_device_array(self):
       )
       np.testing.assert_array_equal(o.data, global_data[o.index])
 
+  @jtu.ignore_warning(category=DeprecationWarning)
   def test_host_local_to_global_replicated(self):
     num_local_devices = jax.local_device_count()
     global_mesh = jax.sharding.Mesh(jax.devices(), axis_names=['x'])
@@ -435,6 +436,7 @@ def test_host_local_to_global_replicated(self):
     # Array is accessible on every host.
     np.testing.assert_array_equal(out, local_input_data)
 
+  @jtu.ignore_warning(category=DeprecationWarning)
   def test_host_local_to_global_locally_replicated(self):
     # Make an array which is locally replicated but sharded across hosts.
     num_processes = jax.process_count()