Repository: jax-ml/jax
Pull request #25843 (2025-01-10 21:40:51)

[sharding_in_types] Introduce `auto_mode`, `user_mode`, `auto_mode_ctx` and `user_mode_ctx` as **private** APIs to make writing auto/user sharding in types code way easier and noise-free.
[sharding_in_types] Introduce `auto_mode`, `user_mode`, `auto_mode_ctx` and `user_mode_ctx` as **private** APIs to make writing auto/user sharding in types code way easier and noise-free.

These can be made public in the future under different names.



Patch:
diff --git a/jax/_src/core.py b/jax/_src/core.py
index c33f28824..22efe720a 100644
--- a/jax/_src/core.py
+++ b/jax/_src/core.py
@@ -1653,19 +1653,21 @@ def _invalid_shape_error(shape: Shape, context: str=""):
 
 # TODO(yashkatariya): Only works with User/Auto. Generalize it to work with
 # Collective too.
-def _maybe_modify_sharding(sharding):
-  if mesh_lib.AxisTypes.Auto not in sharding.mesh.axis_types:
-    return sharding
-
-  new_spec = []
-  for s in sharding.spec:
+def modify_spec_for_auto(spec, mesh) -> P:
+  new_spec = []  # type: ignore
+  for s in spec:
     if s is None:
       new_spec.append(s)
     else:
       temp_s = s[0] if isinstance(s, tuple) else s
       new_spec.append(
-          None
-          if sharding.mesh._name_to_type[temp_s] == mesh_lib.AxisTypes.Auto else s)
+          None if mesh._name_to_type[temp_s] == mesh_lib.AxisTypes.Auto else s)
+  return P(*new_spec)
+
+def _maybe_modify_sharding(sharding):
+  if mesh_lib.AxisTypes.Auto not in sharding.mesh.axis_types:
+    return sharding
+  new_spec = modify_spec_for_auto(sharding.spec, sharding.mesh)
   return sharding.with_spec(new_spec)
 
 
diff --git a/jax/_src/mesh.py b/jax/_src/mesh.py
index 34f485684..9f68a3f8f 100644
--- a/jax/_src/mesh.py
+++ b/jax/_src/mesh.py
@@ -111,14 +111,15 @@ class AxisTypes(enum.Enum):
     return self.name
 
 def axis_names_to_types(axis_types) -> dict[str, AxisTypes]:
-  d = {}
-  for t, names in axis_types.items():
-    if isinstance(names, tuple):
-      for n in names:
-        d[n] = t
-    else:
-      d[names] = t
-  return d
+  return {n: t for t, names in axis_types.items()
+          for n in ((names,) if not isinstance(names, tuple) else names)}
+
+def axis_types_to_names(name_to_type: dict[str, AxisTypes]):
+  d = collections.defaultdict(list)
+  for n, t in name_to_type.items():
+    d[t].append(n)
+  return {t: ns[0] if len(ns) == 1 else tuple(ns) for t, ns in d.items()}
+
 
 _mesh_object_dict = {}  # type: ignore
 
@@ -350,9 +351,6 @@ class Mesh(contextlib.ContextDecorator):
   def abstract_mesh(self):
     return AbstractMesh(self.shape_tuple, axis_types=self.axis_types)
 
-  def with_axis_types(self, new_axis_types) -> Mesh:
-    return Mesh(self.devices, self.axis_names, axis_types=new_axis_types)
-
   @functools.cached_property
   def _are_all_axes_collective(self) -> bool:
     return all(t == AxisTypes.Collective for t in self.axis_types.keys())
@@ -447,7 +445,11 @@ class AbstractMesh:
   def empty(self):
     return self.size == 0
 
-  def with_axis_types(self, new_axis_types) -> AbstractMesh:
+  def update_axis_types(self, new_axis_types) -> AbstractMesh:
+    # dict(self._name_to_type) will copy it.
+    updated_name_to_type = dict(self._name_to_type)
+    updated_name_to_type.update(axis_names_to_types(new_axis_types))
+    new_axis_types = axis_types_to_names(updated_name_to_type)
     return AbstractMesh(self.shape_tuple, axis_types=new_axis_types)
 
   @functools.cached_property
diff --git a/jax/_src/pjit.py b/jax/_src/pjit.py
index 64c606c5c..20d913449 100644
--- a/jax/_src/pjit.py
+++ b/jax/_src/pjit.py
@@ -16,6 +16,7 @@ from __future__ import annotations
 
 from collections import defaultdict
 from collections.abc import Callable, Sequence, Iterable
+import contextlib
 import dataclasses
 from functools import partial
 import inspect
@@ -2738,6 +2739,55 @@ mlir.register_lowering(sharding_cast_p, _sharding_cast_hlo_lowering)
 # batching.fancy_primitive_batchers[sharding_cast_p] = _sharding_cast_batcher
 # batching.skippable_batchers[sharding_cast_p] = lambda _: ()
 
+# -------------------- auto and user mode -------------------------
+
+def _get_new_mesh(axes: str | tuple[str, ...], axis_type: mesh_lib.AxisTypes):
+  if not isinstance(axes, tuple):
+    axes = (axes,)
+  cur_mesh = mesh_lib.get_abstract_mesh()
+  for a in axes:
+    if cur_mesh._name_to_type[a] == axis_type:  # type: ignore
+      raise ValueError(f'Axes {a} cannot be casted to type {axis_type} since '
+                       f'it already is of type {axis_type}.')
+  new_mesh = cur_mesh.update_axis_types({axis_type: axes})  # type: ignore
+  return new_mesh
+
+def auto_mode(fun, *, axes: str | tuple[str, ...], out_specs):
+  new_mesh = _get_new_mesh(axes, mesh_lib.AxisTypes.Auto)
+  def decorator(*args, **kwargs):
+    with mesh_lib.set_abstract_mesh(new_mesh):
+      in_specs = tree_map(lambda a: core.modify_spec_for_auto(
+          a.sharding.spec, new_mesh), args)
+      args = sharding_cast(args, in_specs)
+      out = fun(*args, **kwargs)
+    return sharding_cast(out, out_specs)
+  return decorator
+
+
+@contextlib.contextmanager
+def auto_mode_ctx(axes: str | tuple[str, ...]):
+  new_mesh = _get_new_mesh(axes, mesh_lib.AxisTypes.Auto)
+  with mesh_lib.set_abstract_mesh(new_mesh):
+    yield
+
+
+def user_mode(fun, *, axes: str | tuple[str, ...], in_specs):
+  new_mesh = _get_new_mesh(axes, mesh_lib.AxisTypes.User)
+  def decorator(*args, **kwargs):
+    with mesh_lib.set_abstract_mesh(new_mesh):
+      args = sharding_cast(args, in_specs)
+      out = fun(*args, **kwargs)
+    out_specs = tree_map(lambda o: core.modify_spec_for_auto(
+        o.sharding.spec, mesh_lib.get_abstract_mesh()), out)
+    return sharding_cast(out, out_specs)
+  return decorator
+
+@contextlib.contextmanager
+def user_mode_ctx(axes: str | tuple[str, ...]):
+  new_mesh = _get_new_mesh(axes, mesh_lib.AxisTypes.User)
+  with mesh_lib.set_abstract_mesh(new_mesh):
+    yield
+
 # -------------------- helpers --------------------
 
 def get_unconstrained_dims(sharding: NamedSharding):


Test patch:
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index 44768cbed..95e639e17 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -51,7 +51,8 @@ from jax._src import sharding_impls
 from jax._src.sharding_impls import (
     AUTO, UNSPECIFIED, NamedSharding, GSPMDSharding, PositionalSharding,
     SingleDeviceSharding, parse_flatten_op_sharding)
-from jax._src.pjit import pjit, sharding_cast
+from jax._src.pjit import (pjit, sharding_cast, auto_mode, user_mode,
+                           auto_mode_ctx, user_mode_ctx)
 from jax._src import mesh as mesh_lib
 from jax._src.mesh import set_abstract_mesh, get_abstract_mesh, AxisTypes
 from jax._src.interpreters import pxla
@@ -5672,9 +5673,7 @@ class ShardingInTypesTest(jtu.JaxTestCase):
     @jax.jit
     def f(x):
       y = x * 2
-      auto_mesh = mesh_lib.get_abstract_mesh().with_axis_types(
-          {mesh_lib.AxisTypes.Auto: ('x', 'y')})
-      with mesh_lib.set_abstract_mesh(auto_mesh):
+      with auto_mode_ctx(axes=('x', 'y')):
         y = sharding_cast(y, P(None, None))
         self.assertEqual(y.sharding.spec, P(None, None))
         z = jnp.sin(y)
@@ -5702,10 +5701,8 @@ class ShardingInTypesTest(jtu.JaxTestCase):
     @jax.jit
     def f(x):
       y = x * 2
-      user_mesh = mesh_lib.get_abstract_mesh().with_axis_types(
-          {mesh_lib.AxisTypes.User: ('x', 'y')})
-      y = sharding_cast(y, NamedSharding(user_mesh, P(None, 'y')))
-      with mesh_lib.set_abstract_mesh(user_mesh):
+      with user_mode_ctx(axes=('x', 'y')):
+        y = sharding_cast(y, P(None, 'y'))
         self.assertEqual(y.sharding.spec, P(None, 'y'))
         z = jnp.sin(y)
         self.assertEqual(z.sharding.spec, P(None, 'y'))
@@ -5730,9 +5727,7 @@ class ShardingInTypesTest(jtu.JaxTestCase):
     @jax.jit
     def f(x):
       y = x * 2
-      mix_mesh = mesh_lib.get_abstract_mesh().with_axis_types(
-          {mesh_lib.AxisTypes.Auto: 'x', mesh_lib.AxisTypes.User: 'y'})
-      with mesh_lib.set_abstract_mesh(mix_mesh):
+      with auto_mode_ctx('x'):
         y = sharding_cast(y, P(None, 'y'))
         self.assertEqual(y.sharding.spec, P(None, 'y'))
         z = jnp.sin(y)
@@ -5759,9 +5754,7 @@ class ShardingInTypesTest(jtu.JaxTestCase):
     @jax.jit
     def f(x, y):
       x = x * 2
-      mix_mesh = mesh_lib.get_abstract_mesh().with_axis_types(
-          {mesh_lib.AxisTypes.Auto: 'x', mesh_lib.AxisTypes.User: 'y'})
-      with mesh_lib.set_abstract_mesh(mix_mesh):
+      with auto_mode_ctx('x'):
         z = x @ y
       return z
 
@@ -5824,7 +5817,7 @@ class ShardingInTypesTest(jtu.JaxTestCase):
 
     @jax.jit
     def f(x):
-      auto_mesh = get_abstract_mesh().with_axis_types({AxisTypes.Auto: 'x'})
+      auto_mesh = get_abstract_mesh().update_axis_types({AxisTypes.Auto: 'x'})
       with set_abstract_mesh(auto_mesh):
         x = sharding_cast(x, P(None, None))
         return x
@@ -5873,7 +5866,7 @@ class ShardingInTypesTest(jtu.JaxTestCase):
 
     @jax.jit
     def g(x, y):
-      with mesh_lib.set_abstract_mesh(auto_mesh):
+      with auto_mode_ctx('x'):
         out = jnp.einsum('xy,yz->xz', x, y, out_type=P('x', None))
       return out
 
@@ -5906,6 +5899,117 @@ class ShardingInTypesTest(jtu.JaxTestCase):
     else:
       self.assertTrue(lowered_text.count("unspecified_dims=[1,2]") == 3)
 
+  @jtu.with_user_mesh((2, 2), ('x', 'y'))
+  def test_auto_mode_mix(self, mesh):
+    np_inp = np.arange(16.).reshape(8, 2)
+    s = NamedSharding(mesh, P('x', 'y'))
+    arr = jax.device_put(np_inp, s)
+
+    @partial(auto_mode, axes='x', out_specs=P('x', None))
+    def h(y):
+      self.assertEqual(y.sharding.spec, P(None, 'y'))
+      z = jnp.sin(y)
+      self.assertEqual(z.sharding.spec, P(None, 'y'))
+      a = z @ z.T
+      self.assertEqual(a.sharding.spec, P(None, None))
+      return a
+
+    @jax.jit
+    def g(x):
+      y = x * 2
+      a = h(y)
+      self.assertEqual(a.sharding.spec, P('x', None))
+      return a
+
+    out = g(arr)
+    self.assertEqual(out.sharding, NamedSharding(mesh, P('x', None)))
+
+    jaxpr = g.trace(arr).jaxpr
+    out2 = core.jaxpr_as_fun(jaxpr)(arr)
+    self.assertEqual(out2[0].sharding, NamedSharding(mesh, P('x', None)))
+
+  @jtu.with_user_mesh((2, 2), ('x', 'y'),
+                      axis_types={mesh_lib.AxisTypes.Auto: ('x', 'y')})
+  def test_full_user_mode(self, mesh):
+    np_inp = np.arange(16.).reshape(8, 2)
+    s = NamedSharding(mesh, P('x', 'y'))
+    arr = jax.device_put(np_inp, s)
+
+    @partial(user_mode, axes=('x', 'y'), in_specs=P('x', 'y'))
+    def h(y):
+      self.assertEqual(y.sharding.spec, P('x', 'y'))
+      z = jnp.sin(y)
+      self.assertEqual(z.sharding.spec, P('x', 'y'))
+      a = jnp.einsum('ab,bc->ac', z, z.T, out_type=P('x', None))
+      self.assertEqual(a.sharding.spec, P('x', None))
+      return a
+
+    @jax.jit
+    def f(x):
+      y = x * 2
+      a = h(y)
+      self.assertEqual(a.sharding.spec, P(None, None))
+      return a
+
+    out = f(arr)
+    self.assertEqual(out.sharding, NamedSharding(mesh, P('x')))
+
+    jaxpr = f.trace(arr).jaxpr
+    core.jaxpr_as_fun(jaxpr)(arr)  # doesn't crash
+
+  @jtu.with_user_mesh((2, 2), ('x', 'y'),
+                      axis_types={mesh_lib.AxisTypes.User: 'x',
+                                  mesh_lib.AxisTypes.Auto: 'y'})
+  def test_mix_to_full_user_mode(self, mesh):
+    np_inp = np.arange(16.).reshape(8, 2)
+    s = NamedSharding(mesh, P('x', 'y'))
+    arr = jax.device_put(np_inp, s)
+
+    @partial(user_mode, axes='y', in_specs=P('x', 'y'))
+    def h(y):
+      self.assertEqual(y.sharding.spec, P('x', 'y'))
+      z = jnp.sin(y)
+      self.assertEqual(z.sharding.spec, P('x', 'y'))
+      a = jnp.einsum('ab,bc->ac', z, z.T, out_type=P('x', 'y'))
+      self.assertEqual(a.sharding.spec, P('x', 'y'))
+      return a
+
+    @jax.jit
+    def f(x):
+      y = x * 2
+      a = h(y)
+      self.assertEqual(a.sharding.spec, P('x', None))
+      return a
+
+    out = f(arr)
+    self.assertEqual(out.sharding, NamedSharding(mesh, P('x', 'y')))
+
+  @jtu.with_user_mesh((2, 2), ('x', 'y'),
+                      axis_types={mesh_lib.AxisTypes.Auto: ('x', 'y')})
+  def test_full_auto_to_partial_user(self, mesh):
+    np_inp = np.arange(16.).reshape(8, 2)
+    s = NamedSharding(mesh, P('x', 'y'))
+    arr = jax.device_put(np_inp, s)
+
+    @partial(user_mode, axes='y', in_specs=P(None, 'y'))
+    def h(y):
+      self.assertEqual(y.sharding.spec, P(None, 'y'))
+      z = jnp.sin(y)
+      self.assertEqual(z.sharding.spec, P(None, 'y'))
+      a = jnp.einsum('ab,bc->ac', z, z.T, out_type=P(None, 'y'))
+      self.assertEqual(a.sharding.spec, P(None, 'y'))
+      return a
+
+    @jax.jit
+    def f(x):
+      y = x * 2
+      a = h(y)
+      self.assertEqual(a.sharding.spec, P(None, None))
+      return a
+
+    out = f(arr)
+    self.assertEqual(out.sharding, NamedSharding(mesh, P('x', 'y')))
+
 
 @jtu.pytest_mark_if_available('multiaccelerator')
 class PJitErrorTest(jtu.JaxTestCase):