Fix bugs wrt partial-auto when there are multiple levels of nesting.

The changes are:

* If the mesh passed to shard_map doesn't match the context mesh (if present), error out
* Whenever we trace a jaxpr in shard_map:
  * the avals passed via `_shard_aval` should union the current manual axes on mesh with the newly manual axes specified on shard_map's `axis_name` argument
  * The mesh we enter into when in `use_abstract_mesh` should also union the current manual axes with the newly manual axes

PiperOrigin-RevId: 750774376

Author: yashk2810

jax/_src/debugging.py

@@ -141,7 +141,7 @@ def f():
     return jax.lax.cond(idx == 0,
                         lambda: debug_callback_p.bind(*args, **params),
                         lambda: [])
-  return jax.shard_map(f, mesh=axis_context.mesh, in_specs=(), out_specs=[])()
+  return jax.shard_map(f, in_specs=(), out_specs=[])()
 
 def debug_callback_lowering(ctx, *args, effect, partitioned, callback, **params):
   axis_context = ctx.module_context.axis_context

jax/_src/lax/parallel.py

@@ -1925,8 +1925,8 @@ def _build_axis_index_lowering_hlo(ctx, axis_name, axis_env):
     def f():
       return axis_index_p.bind(axis_name=axis_name)
     return mlir.lower_fun(
-        lambda: [jax.shard_map(f, mesh=axis_context.mesh, check_vma=False,
-                               in_specs=(), out_specs=P())()])(ctx)[0]
+        lambda: [jax.shard_map(f, check_vma=False, in_specs=(),
+                               out_specs=P())()])(ctx)[0]
 
   nreplicas = axis_env.nreps // math.prod(axis_env.sizes)
   div = mlir.ir_constant(

jax/_src/shard_map.py

@@ -128,7 +128,8 @@ def shard_map(f=None, /, *, out_specs: Specs, axis_names: Set[AxisName] = set(),
 
 def _shard_map(f: Callable, *, mesh: Mesh | AbstractMesh | None,
                in_specs: Specs, out_specs: Specs | Callable[[], Specs],
-               axis_names: Set[AxisName], check_vma: bool):
+               axis_names: Set[AxisName], check_vma: bool,
+               _skip_mesh_check: bool = False):
   if not callable(f):
     raise TypeError("shard_map requires a callable for its first argument, "
                     f"but got {f} of type {type(f)}.")
@@ -140,6 +141,14 @@ def _shard_map(f: Callable, *, mesh: Mesh | AbstractMesh | None,
           "The context mesh cannot be empty. Either use"
           " `jax.sharding.use_mesh(mesh)` to enter into a mesh context or pass"
           " a mesh to `shard_map` via the `mesh` keyword argument.")
+  else:
+    ctx_mesh = get_abstract_mesh()
+    if (not _skip_mesh_check and not ctx_mesh.empty and
+        mesh.abstract_mesh != ctx_mesh):
+      raise ValueError(
+          f"The context mesh {ctx_mesh} should match the mesh passed to"
+          f" shard_map {mesh}")
+
   if not isinstance(mesh, (Mesh, AbstractMesh)):
     raise TypeError("shard_map requires a `jax.sharding.Mesh` or a "
                     "`jax.sharding.AbstractMesh` instance for its "
@@ -540,7 +549,7 @@ def _as_manual_mesh(mesh, manual_axes: frozenset):
     if cur_mesh._name_to_type[a] == AxisType.Auto:
       auto_axes.add(a)
     else:
-      assert cur_mesh._name_to_type[a] == AxisType.Explicit
+      assert cur_mesh._name_to_type[a] == AxisType.Explicit, cur_mesh._name_to_type[a]
       explicit_axes.add(a)
 
   new_axis_types = []
@@ -558,7 +567,7 @@ def _as_manual_mesh(mesh, manual_axes: frozenset):
 
 def _extend_axis_env(mesh, manual_axes):
   return core.extend_axis_env_nd([(k, v) for k, v in mesh.shape.items()
-                                 if k in manual_axes])
+                                  if k in manual_axes])
 
 def _shard_map_staging(
     trace: pe.DynamicJaxprTrace, prim: core.Primitive, f: lu.WrappedFun,
@@ -571,11 +580,11 @@ def _shard_map_staging(
   source_info = source_info_util.current()
   to_jaxpr_tracer = partial(trace.to_jaxpr_tracer, source_info=source_info)
   in_tracers = map(to_jaxpr_tracer, in_tracers)
+  inner_mesh = _as_manual_mesh(mesh, manual_axes | set(mesh.manual_axes))
   in_avals = [t.aval for t in in_tracers]
   in_avals_ = map(partial(_shard_aval, mesh, manual_axes, check_vma), in_names,
                   in_avals)
-  manual_mesh = _as_manual_mesh(mesh, manual_axes)
-  with (_extend_axis_env(mesh, manual_axes), use_abstract_mesh(manual_mesh),
+  with (_extend_axis_env(mesh, manual_axes), use_abstract_mesh(inner_mesh),
         config._check_vma(check_vma)):
     jaxpr, out_avals_, consts, () = pe.trace_to_jaxpr_dynamic(f, in_avals_)
   _check_names(out_names_thunk(), out_avals_)
@@ -590,7 +599,7 @@ def _shard_map_staging(
   constvars = map(trace.getvar, map(to_jaxpr_tracer, consts))
   outvars = map(trace.makevar, out_tracers)
   in_names_staged = ({},) * len(consts) + tuple(in_names)  # type: ignore
-  with (_extend_axis_env(mesh, manual_axes), use_abstract_mesh(manual_mesh),
+  with (_extend_axis_env(mesh, manual_axes), use_abstract_mesh(inner_mesh),
         config._check_vma(check_vma)):
     jaxpr = pe.convert_constvars_jaxpr(jaxpr)
   params = dict(mesh=mesh, in_names=in_names_staged,
@@ -629,10 +638,11 @@ def _shard_shaped_array(mesh: Mesh, manual_axes: frozenset, check_vma,
   assert isinstance(aval, core.ShapedArray)
   new_shape = tuple(sz // prod(mesh.shape[n] for n in names.get(i, ()))
                     for i, sz in enumerate(aval.shape))
-  manual_mesh = _as_manual_mesh(mesh, manual_axes)
+  manual_mesh = _as_manual_mesh(mesh, manual_axes | set(mesh.manual_axes))
   new_sharding = NamedSharding(manual_mesh, aval.sharding.spec)
   vma = (frozenset({n for ns in names.values() for n in ns})
          if check_vma else frozenset())
+  vma = vma | aval.vma
   return aval.update(shape=new_shape, sharding=new_sharding, vma=vma)
 core.shard_aval_handlers[core.ShapedArray] = _shard_shaped_array
 
@@ -695,7 +705,7 @@ def _shard_map_typecheck(_, *in_atoms, jaxpr, mesh, in_names, out_names,
 
 
 def _valid_repeats(mesh: Mesh, vma: Set[AxisName], names: AxisNames) -> bool:
-  um = set(_unmentioned(mesh, names))
+  um = set(_unmentioned(mesh, names)) - set(mesh.manual_axes)
   if any(u in vma for u in um):
     return False
   return True
@@ -808,8 +818,10 @@ def _xla_shard(ctx: mlir.LoweringRuleContext, mesh, manual_axes, names,
                  if len(manual_axes) < len(mesh.axis_names) else set())
   sx = mlir.wrap_with_sharding_op(ctx, x, aval_in, shard_proto,
                                   unspecified_dims=unspecified)
-  manual_proto = pxla.manual_proto(aval_in, manual_axes, mesh)
-  return mlir.wrap_with_full_to_shard_op(ctx, sx, aval_out, manual_proto, unspecified)
+  manual_proto = pxla.manual_proto(
+      aval_in, manual_axes | set(mesh.manual_axes), mesh)
+  return mlir.wrap_with_full_to_shard_op(ctx, sx, aval_out, manual_proto,
+                                         unspecified)
 
 def _xla_unshard(ctx: mlir.LoweringRuleContext, mesh, manual_axes, names,
                  aval_in, aval_out, x):
@@ -824,8 +836,10 @@ def _xla_unshard(ctx: mlir.LoweringRuleContext, mesh, manual_axes, names,
                  if len(manual_axes) < len(mesh.axis_names) else set())
   if dtypes.issubdtype(aval_in.dtype, dtypes.extended):
     aval_in = core.physical_aval(aval_in)
-  manual_proto = pxla.manual_proto(aval_in, manual_axes, mesh)
-  sx = mlir.wrap_with_sharding_op(ctx, x, aval_in, manual_proto, unspecified_dims=unspecified)
+  manual_proto = pxla.manual_proto(
+      aval_in, manual_axes | set(mesh.manual_axes), mesh)
+  sx = mlir.wrap_with_sharding_op(ctx, x, aval_in, manual_proto,
+                                  unspecified_dims=unspecified)
   shard_proto = ns._to_xla_hlo_sharding(aval_out.ndim).to_proto()
   return mlir.wrap_with_shard_to_full_op(ctx, sx, aval_out, shard_proto,
                                          unspecified)
@@ -894,9 +908,9 @@ def _shard_map_impl(trace, prim, fun, args, *, mesh, in_names, out_names_thunk,
 def _run_shmap(f, mesh, manual_axes, args, vmas, check_vma, context_mesh):
   trace = ShardMapTrace(mesh, manual_axes, check_vma, context_mesh)
   in_tracers = map(partial(ShardMapTracer, trace), vmas, args)
-  manual_mesh = _as_manual_mesh(mesh, manual_axes)
+  inner_mesh = _as_manual_mesh(mesh, manual_axes | set(mesh.manual_axes))
   with (core.set_current_trace(trace), _extend_axis_env(mesh, manual_axes),
-        use_abstract_mesh(manual_mesh), config._check_vma(check_vma)):
+        use_abstract_mesh(inner_mesh), config._check_vma(check_vma)):
     ans = f.call_wrapped(*in_tracers)
     outs, out_vma = unzip2(map(trace.to_val_vma_pair, ans))
   return outs, out_vma
@@ -1318,7 +1332,7 @@ def fwd_out_names_thunk():
   args_to_promote = [getattr(aval, 'shape', ()) == () and f1 is None and f2 is None
                      for aval, f1, f2 in zip(res_avals, in_fwd, out_fwd)]
   with (_extend_axis_env(mesh, manual_axes),
-        use_abstract_mesh(_as_manual_mesh(mesh, manual_axes)),
+        use_abstract_mesh(_as_manual_mesh(mesh, manual_axes | set(mesh.manual_axes))),
         config._check_vma(check_vma)):
     lin_jaxpr = _promote_scalar_residuals_jaxpr(lin_jaxpr, args_to_promote)
   out_names = out_names_thunk()
@@ -1483,7 +1497,7 @@ def _partial_eval_jaxpr_custom_rule(
   jaxpr, mesh = eqn.params['jaxpr'], eqn.params['mesh']
   check_vma, manual_axes = eqn.params['check_vma'], eqn.params['manual_axes']
   with (_extend_axis_env(mesh, manual_axes), config._check_vma(check_vma),
-        use_abstract_mesh(_as_manual_mesh(mesh, manual_axes))):
+        use_abstract_mesh(_as_manual_mesh(mesh, manual_axes | set(mesh.manual_axes)))):
     jaxpr_known, jaxpr_staged, unks_out, inst_out, num_res = \
         pe.partial_eval_jaxpr_custom(jaxpr, unks_in, inst_in, False, False, saveable)
   num_out_primals = len(jaxpr_known.outvars) - num_res
@@ -1494,7 +1508,7 @@ def _partial_eval_jaxpr_custom_rule(
   which = [f1 is None and f2 is None for f1, f2 in zip(in_fwd, out_fwd)]
   mesh = eqn.params['mesh']
   with (_extend_axis_env(mesh, manual_axes),
-        use_abstract_mesh(_as_manual_mesh(mesh, manual_axes)),
+        use_abstract_mesh(_as_manual_mesh(mesh, manual_axes | set(mesh.manual_axes))),
         config._check_vma(check_vma)):
     jaxpr_known = pe.prune_jaxpr_outputs(jaxpr_known, [True] * num_out_primals + which)
     jaxpr_known, jaxpr_staged = _add_reshapes(which, jaxpr_known, jaxpr_staged)

jax/experimental/shard_map.py

@@ -77,6 +77,6 @@ def shard_map(
   .. _SPMD multi-device parallelism with shard_map: https://docs.jax.dev/en/latest/notebooks/shard_map.html
   """
   axis_names = frozenset(mesh.axis_names) - auto
-  return jshmap.shard_map(
+  return jshmap._shard_map(
       f, mesh=mesh, in_specs=in_specs, out_specs=out_specs,
-      check_vma=check_rep, axis_names=axis_names)
+      check_vma=check_rep, axis_names=axis_names, _skip_mesh_check=True)

tests/shard_map_test.py

@@ -438,19 +438,17 @@ def test_replication_checker_jit(self):
     mesh = jtu.create_mesh((2, 2), ('x', 'y'))
     x = np.arange(8 * 8.).reshape(8, 8)
 
-    def f(x):
-      return 2 * x
     def g(x):
-      return shard_map(f, mesh=mesh, in_specs=(P('x', 'y'),), out_specs=P(None, 'y'))(x)
+      return shard_map(lambda x: x * 2, mesh=mesh, in_specs=P('x', 'y'),
+                       out_specs=P(None, 'y'))(x)
 
     with self.assertRaisesRegex(ValueError, 'statically inferred'):
       jax.jit(g)(x)
 
-    def f2(x):
-      return jax.lax.psum(x, 'x')
     def g2(x):
-      return shard_map(f2, mesh=mesh, in_specs=(P('x', 'y'),), out_specs=P(None, 'y'))(x)
-    _ = jax.jit(g2)(x)  # doesn't crash
+      return shard_map(lambda x: jax.lax.psum(x, 'x'), mesh=mesh,
+                       in_specs=P('x', 'y'), out_specs=P(None, 'y'))(x)
+    jax.jit(g2)(x)  # doesn't crash
 
   def test_process_env_traces(self):
     mesh = Mesh(np.array(jax.devices()[:4]), ('x',))
@@ -2242,16 +2240,54 @@ def g(x):
       return x * x
 
     def h(x):
-      return shard_map(g, mesh=mesh, in_specs=P(None, 'j'), out_specs=P(None, 'j'))(x)
+      return shard_map(g, in_specs=P(None, 'j'), out_specs=P(None, 'j'))(x)
 
     @jax.jit
     def f(x):
-      return shard_map(h, mesh=mesh, in_specs=P('i', None), out_specs=P('i', None),
+      return shard_map(h, in_specs=P('i', None), out_specs=P('i', None),
                        check_vma=False, axis_names=frozenset({'i'}))(x)
 
     v = jnp.arange(32.).reshape(4, 8)
     v = jax.device_put(v, jax.sharding.NamedSharding(mesh, P('i', 'j')))
-    self.assertAllClose(v*v, f(v), check_dtypes=False)
+    with jax.sharding.use_mesh(mesh):
+      self.assertAllClose(v*v, f(v), check_dtypes=False)
+
+  @parameterized.named_parameters(
+      ('0', 'x', 'y', {'x'}, {'x', 'y'}),
+      ('1', None, 'y', frozenset(), {'y'}),
+      ('2', 'x', None, {'x'}, {'x'}),
+      ('3', None, None, frozenset(), frozenset()),
+  )
+  def test_nested_partial_auto_1d(self, dim1, dim2, outer_vma, inner_vma):
+    mesh = jtu.create_mesh((2, 2, 2), ('x', 'y', 'z'))
+    np_inp = np.arange(32.).reshape(4, 8)
+    arr = jax.device_put(np_inp, NamedSharding(mesh, P(dim1, dim2)))
+
+    def g(x):
+      self.assertEqual(get_abstract_mesh().manual_axes, ('x', 'y'))
+      self.assertEqual(get_abstract_mesh().auto_axes, ('z',))
+      self.assertEqual(x.aval.vma, inner_vma)
+      out = x * x
+      self.assertEqual(out.aval.vma, inner_vma)
+      return out
+
+    def h(x):
+      self.assertEqual(get_abstract_mesh().manual_axes, ('x',))
+      self.assertEqual(get_abstract_mesh().auto_axes, ('y', 'z'))
+      self.assertEqual(x.aval.vma, outer_vma)
+      out = shard_map(g, in_specs=P(None, dim2),
+                      out_specs=P(None, dim2), axis_names={'y'})(x)
+      self.assertEqual(out.aval.vma, outer_vma)
+      return out
+
+    @jax.jit
+    def f(x):
+      return shard_map(h, in_specs=P(dim1, None),
+                       out_specs=P(dim1, None), axis_names={'x'})(x)
+
+    with jax.sharding.use_mesh(mesh):
+      out = f(arr)
+      self.assertArraysEqual(out, np_inp * np_inp)
 
   def test_grad_nested_partial_auto(self):
     mesh = jtu.create_mesh((2, 2), ('i', 'j'))
@@ -2262,22 +2298,19 @@ def g(x):
 
     def h(x):
       # auto: 'j', manual: 'i'
-      return shard_map(g, mesh=mesh,
-                       in_specs=P(None, 'j'),
-                       out_specs=P(None, 'j'))(x)
+      return shard_map(g, in_specs=P(None, 'j'), out_specs=P(None, 'j'))(x)
 
     @jax.jit
     def f(x):
       # auto: 'i', 'j'
-      return shard_map(h, mesh=mesh,
-                       in_specs=P('i', None),
-                       out_specs=P('i', None),
-                       check_vma=False,
-                       axis_names=frozenset({'i'}))(x).sum()
+      return shard_map(h, in_specs=P('i', None), out_specs=P('i', None),
+                       check_vma=False, axis_names=frozenset({'i'}))(x).sum()
 
     v = jnp.arange(32.).reshape(4, 8)
     v = jax.device_put(v, jax.sharding.NamedSharding(mesh, P('i', 'j')))
-    self.assertAllClose(v*2, jax.grad(f)(v), check_dtypes=False)
+    with jax.sharding.use_mesh(mesh):
+      out = jax.grad(f)(v)
+      self.assertAllClose(out, v * 2, check_dtypes=False)
 
   def test_grad_nested_partial_auto_with_residuals(self):
     mesh = jtu.create_mesh((2, 2), ('i', 'j'))
@@ -2286,21 +2319,18 @@ def g(x):
       return x * x * x
 
     def h(x):
-      return shard_map(g, mesh=mesh,
-                    in_specs=P(None, 'j'),
-                    out_specs=P(None, 'j'))(x)
+      return shard_map(g, in_specs=P(None, 'j'), out_specs=P(None, 'j'))(x)
 
     @jax.jit
     def f(x):
-      return shard_map(h, mesh=mesh,
-                    in_specs=P('i', None),
-                    out_specs=P('i', None),
-                    check_vma=False,
-                    axis_names=frozenset({'i'}))(x).sum()
+      return shard_map(h, in_specs=P('i', None), out_specs=P('i', None),
+                       check_vma=False, axis_names=frozenset({'i'}))(x).sum()
 
     v = jnp.arange(32.).reshape(4, 8)
     v = jax.device_put(v, jax.sharding.NamedSharding(mesh, P('i', 'j')))
-    self.assertAllClose(v*v*3, jax.grad(f)(v), check_dtypes=False)
+    with jax.sharding.use_mesh(mesh):
+      out = jax.grad(f)(v)
+      self.assertAllClose(out, v * v * 3, check_dtypes=False)
 
   def test_axis_size_1_partial_auto(self):
     mesh = jtu.create_mesh((1, 2, 2), ('i', 'j', 'k'))
@@ -2367,10 +2397,11 @@ def test_partial_auto_axis_index(self):
     @partial(jax.jit, out_shardings=out_sharding)
     def f():
       return shard_map(lambda: jax.lax.axis_index('i').reshape(1,1),
-                       mesh=mesh, in_specs=P('i', None), out_specs=P('i', None),
+                       in_specs=P('i', None), out_specs=P('i', None),
                        check_vma=False, axis_names=frozenset({'i'}))()
 
-    self.assertAllClose(f(), np.arange(4, dtype=np.int32).reshape(-1, 1))
+    with jax.sharding.use_mesh(mesh):
+      self.assertAllClose(f(), np.arange(4, dtype=np.int32).reshape(-1, 1))
 
   def test_partial_auto_axis_index_degenerated_axis(self):
     mesh = jtu.create_mesh((1, 2), ('i', 'j'))
@@ -2432,11 +2463,11 @@ def g(x):
 
     @jax.jit
     def f(x):
-      return shard_map(g,
-                       mesh=mesh, in_specs=P('i'), out_specs=None,
+      return shard_map(g, mesh=mesh, in_specs=P('i'), out_specs=None,
                        check_vma=False, axis_names=frozenset({'i'}))(x)
 
-    y = f(x)  # don't crash
+    with jax.sharding.use_mesh(mesh):
+      f(x)  # don't crash
 
   def test_partial_auto_of_random_keys(self):
     mesh = jtu.create_mesh((4, 2), ('i', 'j'))