Expose jax.sharding.get_mesh() as a way to get the concrete mesh. This is useful for writing libraries. get_mesh() can't be used inside jax.jit just like set_mesh.

PiperOrigin-RevId: 833451082

Author: yashk2810

jax/_src/sharding_impls.py

@@ -1258,6 +1258,14 @@ def __exit__(self, exc_type, exc_value, traceback):
     config.device_context.set_local(self.prev_mesh)
 
 
+def get_mesh() -> mesh_lib.Mesh:
+  if not core.trace_state_clean():
+    raise ValueError(
+        '`get_mesh` can only be used outside of `jax.jit`. Maybe you want'
+        ' `jax.sharding.get_abstract_mesh()`?')
+  return mesh_lib.get_concrete_mesh()
+
+
 @contextlib.contextmanager
 def _internal_use_concrete_mesh(mesh: mesh_lib.Mesh):
   assert isinstance(mesh, mesh_lib.Mesh)

jax/sharding.py

@@ -21,6 +21,7 @@
     SingleDeviceSharding as SingleDeviceSharding,
     PmapSharding as _deprecated_PmapSharding,
     set_mesh as set_mesh,
+    get_mesh as get_mesh,
 )
 from jax._src.partition_spec import (
     PartitionSpec as PartitionSpec,

tests/documentation_coverage_test.py

@@ -71,7 +71,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'],
+  'jax.sharding': ['AbstractDevice', 'AbstractMesh', 'AxisType', 'auto_axes', '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'],
 }

tests/pjit_test.py

@@ -7633,6 +7633,8 @@ def test_set_mesh(self):
       jax.set_mesh(mesh)
       out = reshard(np.arange(8), P('x'))
       self.assertEqual(out.sharding, NamedSharding(mesh, P('x')))
+      out_mesh = jax.sharding.get_mesh()
+      self.assertEqual(out_mesh, mesh)
     finally:
       config.abstract_mesh_context_manager.set_local(
           mesh_lib.empty_abstract_mesh)