[sharding_in_types] Use set_mesh API to trigger sharding_in_types instead of the config option.

PiperOrigin-RevId: 702814257

Author: yashk2810

jax/_src/interpreters/pxla.py

@@ -2193,15 +2193,8 @@ def lower_sharding_computation(
   assert len(out_shardings) == len(out_layouts) == len(global_out_avals), (
       len(out_shardings), len(out_layouts), len(global_out_avals))
 
-  if config.sharding_in_types.value:
-    # TODO(yashkatariya): Thread it via jit path and remove the None check by
-    # making tests go via set_mesh API always.
-    devices_from_context = (
-        None if mesh_lib.device_context.concrete_mesh is None
-        else mesh_lib.device_context.concrete_mesh._flat_devices_tuple)
-  else:
-    devices_from_context = (None if context_mesh is None or context_mesh.empty
-                            else context_mesh._flat_devices_tuple)
+  devices_from_context = (None if context_mesh is None or context_mesh.empty
+                          else context_mesh._flat_devices_tuple)
   # Device assignment across all inputs, outputs and shardings inside jaxpr
   # should be the same.
   unique_intermediate_shardings = util.stable_unique(

jax/_src/pjit.py

@@ -707,9 +707,6 @@ def get_abstract_mesh(in_avals):
           f'Mesh for all inputs should be equal. Got one mesh: {m} and'
           f' another mesh: {a.sharding.mesh}')
     m = a.sharding.mesh  # type: ignore
-  # TODO(yashkatariya): Remove this when mesh context can be set by the user.
-  if m is None:
-    return contextlib.nullcontext()
   assert isinstance(m, AbstractMesh)
   return m
 
@@ -1791,8 +1788,12 @@ def _pjit_lower(
     lowering_platforms: tuple[str, ...] | None,
     lowering_parameters: mlir.LoweringParameters,
     pgle_profiler: profiler.PGLEProfiler | None):
-  mesh, api_name = ((resource_env.physical_mesh, 'pjit')
-                    if resource_env is not None else (None, 'jit'))
+  if config.sharding_in_types.value:
+    mesh = mesh_lib.device_context.concrete_mesh
+    api_name = 'jit'
+  else:
+    mesh, api_name = ((resource_env.physical_mesh, 'pjit')
+                      if resource_env is not None else (None, 'jit'))
   return pxla.lower_sharding_computation(
       jaxpr, api_name, name, in_shardings, out_shardings,
       in_layouts, out_layouts, tuple(donated_invars),

jax/_src/test_util.py

@@ -51,6 +51,7 @@
 from jax._src import pjit as pjit_lib
 from jax._src import stages
 from jax._src import xla_bridge
+from jax._src import mesh as mesh_lib
 from jax._src.cloud_tpu_init import running_in_cloud_tpu_vm
 from jax._src.interpreters import mlir
 from jax._src.interpreters import pxla
@@ -1442,6 +1443,16 @@ def with_and_without_mesh(f):
       ('Mesh', (('x', 2),), (('i', 'x'),))
     ))(with_mesh_from_kwargs(f))
 
+def with_user_mesh(sizes, names):
+  def decorator(fn):
+    def mesh_fn(*args, **kwargs):
+      mesh = create_mesh(sizes, names)
+      with mesh_lib.set_mesh(mesh):
+        return fn(*args, **kwargs, mesh=mesh)
+    return mesh_fn
+  return decorator
+
+
 def create_mesh(mesh_shape, axis_names, iota_order=False):
   size = math.prod(mesh_shape)
   if len(jax.devices()) < size: