[pmap] Add default_pmap_sharding to migrate users away from PmapSharding.default.

We leave this function undocumented because it's not intended for new users, but is public to migrate existing users.

PiperOrigin-RevId: 842759555

Author: danielsuo

jax/_src/sharding_impls.py

@@ -351,6 +351,52 @@ def shard_shape(self, global_shape: Shape) -> Shape:
 PmapSharding.__module__ = 'jax.sharding'
 
 
+def default_pmap_sharding(
+    shape: Shape,
+    sharded_dim: int | None = 0,
+    devices: Sequence[xc.Device] | None = None,
+) -> NamedSharding | PmapSharding:
+  """Creates a NamedSharding equivalent to PmapSharding.default.
+
+  This function provides the same sharding semantics as PmapSharding.default
+  but returns a NamedSharding when jax_pmap_shmap_merge is enabled, which is
+  compatible with the shard_map-based pmap implementation.
+
+  Args:
+    shape: The shape of the input array.
+    sharded_dim: Dimension the input array is sharded on. Defaults to 0.
+      If None, the array is fully replicated.
+    devices: Optional sequence of devices to use. If omitted, uses
+      jax.local_devices().
+
+  Returns:
+    A NamedSharding if jax_pmap_shmap_merge is enabled, otherwise a
+    PmapSharding.
+  """
+  if not config.pmap_shmap_merge.value:
+    return PmapSharding.default(shape, sharded_dim=sharded_dim, devices=devices)
+
+  if sharded_dim is None:
+    if devices is None:
+      raise ValueError("One of sharded_dim or devices must be set.")
+    mesh = mesh_lib.Mesh(np.array(devices), ('_default_pmap_sharding',))
+    return NamedSharding(mesh, PartitionSpec())
+
+  if len(shape) == 0:
+    raise ValueError("shape must be non-empty for sharded_dim != None")
+
+  num_ways_sharded = shape[sharded_dim]
+
+  if devices is None:
+    pmap_devices = np.array(xb.local_devices()[:num_ways_sharded])
+  else:
+    pmap_devices = np.array(devices)
+
+  mesh = mesh_lib.Mesh(pmap_devices, ('_default_pmap_sharding',))
+  spec_list: list[str | None] = [None] * len(shape)
+  spec_list[sharded_dim] = '_default_pmap_sharding'
+  return NamedSharding(mesh, PartitionSpec(*spec_list))
+
 def _unpickle_gspmd_sharding(devices, op_sharding, memory_kind):
   return GSPMDSharding(devices, op_sharding, memory_kind=memory_kind)
 

jax/sharding.py

@@ -20,6 +20,7 @@
     NamedSharding as NamedSharding,
     SingleDeviceSharding as SingleDeviceSharding,
     PmapSharding as _deprecated_PmapSharding,
+    default_pmap_sharding as default_pmap_sharding,
     set_mesh as set_mesh,
     get_mesh as get_mesh,
 )

tests/documentation_coverage_test.py

@@ -68,7 +68,7 @@ def jax_docs_dir() -> str:
   'jax.profiler': ['ProfileData', 'ProfileEvent', 'ProfileOptions', 'ProfilePlane', 'stop_server'],
   'jax.random': ['key_impl', 'random_gamma_p'],
   'jax.scipy.special': ['bessel_jn', 'sph_harm_y'],
-  'jax.sharding': ['AbstractDevice', 'AbstractMesh', 'AxisType', 'auto_axes', 'explicit_axes', 'get_abstract_mesh', 'reshard', 'set_mesh', 'use_abstract_mesh', 'get_mesh'],
+  'jax.sharding': ['AbstractDevice', 'AbstractMesh', 'AxisType', 'auto_axes', 'default_pmap_sharding', 'explicit_axes', 'get_abstract_mesh', 'reshard', 'set_mesh', 'use_abstract_mesh', 'get_mesh'],
   'jax.stages': ['ArgInfo', 'CompilerOptions'],
   'jax.tree_util': ['DictKey', 'FlattenedIndexKey', 'GetAttrKey', 'PyTreeDef', 'SequenceKey', 'default_registry'],
 }