Repository: google/jax
Pull request #22643 (2024-07-24 23:21:24)

Simplify pjit's batching rule now that xmap is deleted. Also do cleanup around adding manual axes under shard_map
Simplify pjit's batching rule now that xmap is deleted. Also do cleanup around adding manual axes under shard_map



Patch:
diff --git a/jax/_src/interpreters/mlir.py b/jax/_src/interpreters/mlir.py
index a6f4fa14b..08eafedb6 100644
--- a/jax/_src/interpreters/mlir.py
+++ b/jax/_src/interpreters/mlir.py
@@ -874,6 +874,21 @@ class LoweringResult(NamedTuple):
 _platforms_with_donation = ["cpu", "cuda", "rocm", "tpu"]
 
 
+def add_manual_axes(axis_ctx: sharding_impls.SPMDAxisContext, sharding, ndim):
+  mesh = axis_ctx.mesh
+  if (isinstance(sharding, sharding_impls.NamedSharding) and
+      sharding.mesh.shape == mesh.shape):
+    return sharding_impls.NamedSharding._from_parsed_pspec(
+        sharding.mesh, sharding._parsed_pspec, memory_kind=sharding.memory_kind,
+        _manual_axes=axis_ctx.manual_axes)
+  else:
+    parsed_pspec = sharding_impls.parse_flatten_op_sharding(
+      sharding._to_xla_hlo_sharding(ndim), mesh)[0]
+    return sharding_impls.NamedSharding._from_parsed_pspec(
+      mesh, parsed_pspec, memory_kind=sharding.memory_kind,
+      _manual_axes=axis_ctx.manual_axes)
+
+
 def _to_physical_op_sharding(
     ctx: ModuleContext,
     aval: core.AbstractValue, sharding: JSharding | None,
@@ -890,17 +905,7 @@ def _to_physical_op_sharding(
   axis_ctx = ctx.axis_context
   if (isinstance(axis_ctx, sharding_impls.SPMDAxisContext) and
       axis_ctx.manual_axes):
-    mesh = axis_ctx.mesh
-    if (isinstance(sharding, sharding_impls.NamedSharding) and
-        sharding.mesh.shape == mesh.shape):
-      sharding = sharding_impls.NamedSharding._from_parsed_pspec(
-          sharding.mesh, sharding._parsed_pspec,
-          _manual_axes=axis_ctx.manual_axes)
-    else:
-      parsed_pspec = sharding_impls.parse_flatten_op_sharding(
-        sharding._to_xla_hlo_sharding(aval.ndim), mesh)[0]
-      sharding = sharding_impls.NamedSharding._from_parsed_pspec(
-        mesh, parsed_pspec, _manual_axes=axis_ctx.manual_axes)
+    sharding = add_manual_axes(axis_ctx, sharding, aval.ndim)
   if config.use_shardy_partitioner.value:
     return sharding._to_sdy_sharding(aval.ndim)  # type: ignore
   return sharding._to_xla_hlo_sharding(aval.ndim).to_proto()  # type: ignore
diff --git a/jax/_src/pjit.py b/jax/_src/pjit.py
index 32f856be5..648ab3168 100644
--- a/jax/_src/pjit.py
+++ b/jax/_src/pjit.py
@@ -1941,20 +1941,15 @@ def _pjit_lowering(ctx, *args, name, jaxpr, in_shardings,
 mlir.register_lowering(pjit_p, _pjit_lowering)
 
 
-def _pjit_batcher(insert_axis, spmd_axis_name,
-                  axis_size, axis_name, main_type,
-                  vals_in, dims_in,
-                  jaxpr, in_shardings, out_shardings, in_layouts, out_layouts,
-                  resource_env, donated_invars, name, keep_unused, inline):
+def _pjit_batcher(spmd_axis_name, axis_size, axis_name, main_type,
+                  vals_in, dims_in, jaxpr, in_shardings, out_shardings,
+                  in_layouts, out_layouts, resource_env, donated_invars, name,
+                  keep_unused, inline):
   segment_lens, dims_in = batching.indirectify_ragged_axes(dims_in)
   new_jaxpr, axes_out = batching.batch_jaxpr2(
       jaxpr, axis_size, dims_in, axis_name=axis_name,
       spmd_axis_name=spmd_axis_name, main_type=main_type)
 
-  # `insert_axis` is set to True only for some `xmap` uses.
-  new_parts = (axis_name,) if insert_axis else (
-      () if spmd_axis_name is None else spmd_axis_name)
-
   if resource_env is not None:
     mesh = resource_env.physical_mesh
   else:
@@ -1962,11 +1957,11 @@ def _pjit_batcher(insert_axis, spmd_axis_name,
 
   # TODO(axch): prepend with Nones (?) to account for new segment_lens inputs
   in_shardings = tuple(
-      _pjit_batcher_for_sharding(i, axis_in, new_parts, mesh, aval.ndim)
+      _pjit_batcher_for_sharding(i, axis_in, spmd_axis_name, mesh, aval.ndim)
       if axis_in is not None else i
       for axis_in, i, aval in zip(dims_in, in_shardings, new_jaxpr.in_avals))
   out_shardings = tuple(
-      _pjit_batcher_for_sharding(o, axis_out, new_parts, mesh, aval.ndim)
+      _pjit_batcher_for_sharding(o, axis_out, spmd_axis_name, mesh, aval.ndim)
       if axis_out is not None else o
       for axis_out, o, aval in zip(axes_out, out_shardings, new_jaxpr.out_avals))
   # TODO(yashkatariya): Figure out layouts should change under vmap.
@@ -1991,16 +1986,16 @@ def _pjit_batcher(insert_axis, spmd_axis_name,
       vals_in, vals_out, axes_out)
   return vals_out, resolved_axes_out
 
-batching.spmd_axis_primitive_batchers[pjit_p] = partial(_pjit_batcher, False)
-batching.axis_primitive_batchers[pjit_p] = partial(_pjit_batcher, False, None)
+batching.spmd_axis_primitive_batchers[pjit_p] = _pjit_batcher
+batching.axis_primitive_batchers[pjit_p] = partial(_pjit_batcher, None)
 
 def _pjit_batcher_for_sharding(
     s: sharding.Sharding | UnspecifiedValue,
-    dim: int, val: tuple[str, ...], mesh, ndim: int):
+    dim: int, spmd_axis_name: tuple[str, ...] | None, mesh, ndim: int):
   if is_unspecified(s):
     return s
   hlo_s = s._to_xla_hlo_sharding(ndim)  # type: ignore
-  if not val:
+  if spmd_axis_name is None:
     if sharding_impls.is_op_sharding_replicated(hlo_s):
       return s
     new_op = hlo_s.to_proto().clone()
@@ -2016,14 +2011,14 @@ def _pjit_batcher_for_sharding(
       mesh = s.mesh
     if mesh is None or mesh.empty:
       raise ValueError(
-          'If you are using xmap or spmd_axis_name parameter of jax.vmap,'
+          'If you are using spmd_axis_name parameter of jax.vmap,'
           ' please make sure to run your jitted function inside the mesh'
           ' context manager. Only `jax.lax.with_sharding_constraint` with'
           ' `jax.sharding.NamedSharding` as an input can be transformed with'
           ' spmd_axis_name batching rules outside of an explicit mesh context'
           f' manager scope{s!r}')
     parsed_pspec = parse_flatten_op_sharding(hlo_s, mesh)[0]
-    parsed_pspec = parsed_pspec.insert_axis_partitions(dim, val)
+    parsed_pspec = parsed_pspec.insert_axis_partitions(dim, spmd_axis_name)
     return NamedSharding._from_parsed_pspec(mesh, parsed_pspec)
 
 
@@ -2498,16 +2493,9 @@ def _sharding_constraint_hlo_lowering(ctx, x_node, *, sharding, layout,
   aval, = ctx.avals_in
   out_aval, = ctx.avals_out
   axis_ctx = ctx.module_context.axis_context
-  # axis_ctx and manual_axes is *only used with xmap* and xmap only works with
-  # NamedSharding. So update the NamedSharding to have the manual axes.
-  if isinstance(axis_ctx, sharding_impls.SPMDAxisContext):
-    mesh = resource_env.physical_mesh
-    if mesh.empty and isinstance(sharding, NamedSharding):
-      mesh = sharding.mesh
-    parsed_pspec = parse_flatten_op_sharding(
-        sharding._to_xla_hlo_sharding(aval.ndim), mesh)[0]
-    sharding = NamedSharding._from_parsed_pspec(
-        mesh, parsed_pspec, _manual_axes=axis_ctx.manual_axes)
+  if (isinstance(axis_ctx, sharding_impls.SPMDAxisContext) and
+      axis_ctx.manual_axes):
+    sharding = mlir.add_manual_axes(axis_ctx, sharding, aval.ndim)
   out = mlir.wrap_with_sharding_op(
       ctx, x_node, out_aval, sharding._to_xla_hlo_sharding(aval.ndim).to_proto(),
       unspecified_dims=unconstrained_dims)
@@ -2519,7 +2507,7 @@ mlir.register_lowering(sharding_constraint_p,
 
 
 def _sharding_constraint_batcher(
-    insert_axis, spmd_axis_name, axis_size, axis_name, main_type, vals_in,
+    spmd_axis_name, axis_size, axis_name, main_type, vals_in,
     dims_in, sharding, layout, resource_env, unconstrained_dims):
   if spmd_axis_name is not None and isinstance(sharding, NamedSharding):
     used = {n for ns in sharding.spec
@@ -2530,16 +2518,13 @@ def _sharding_constraint_batcher(
                        f"{sharding.spec}")
   x, = vals_in
   d, = dims_in
-  # None means unconstrained in ParsedPartitionSpec
-  new_parts = (axis_name,) if insert_axis else (
-      None if spmd_axis_name is None else spmd_axis_name)
-  unconstrained_dims = {ud + (d <= ud) for ud in unconstrained_dims}
 
-  if new_parts is None:
+  unconstrained_dims = {ud + (d <= ud) for ud in unconstrained_dims}
+  if spmd_axis_name is None:
     unconstrained_dims.add(d)
 
   vmapped_sharding = _pjit_batcher_for_sharding(
-      sharding, d, new_parts, resource_env.physical_mesh, x.ndim)
+      sharding, d, spmd_axis_name, resource_env.physical_mesh, x.ndim)
   if unconstrained_dims and isinstance(vmapped_sharding, NamedSharding):
     new_spec = list(vmapped_sharding.spec) + [None] * (x.ndim - len(vmapped_sharding.spec))
     for u in unconstrained_dims:
@@ -2558,10 +2543,9 @@ def _sharding_constraint_batcher(
       resource_env=resource_env,
       unconstrained_dims=unconstrained_dims)
   return y, d
-batching.spmd_axis_primitive_batchers[sharding_constraint_p] = partial(
-    _sharding_constraint_batcher, False)
+batching.spmd_axis_primitive_batchers[sharding_constraint_p] = _sharding_constraint_batcher
 batching.axis_primitive_batchers[sharding_constraint_p] = partial(
-    _sharding_constraint_batcher, False, None)
+    _sharding_constraint_batcher, None)
 
 # -------------------- helpers --------------------
 


Test patch:
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index 9ed2405bc..97ba67bb9 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -3778,7 +3778,7 @@ class ArrayPjitTest(jtu.JaxTestCase):
     arr = jax.device_put(np.arange(8), s)
     with self.assertRaisesRegex(
         ValueError,
-        'If you are using xmap or spmd_axis_name parameter of jax.vmap, please'
+        'If you are using spmd_axis_name parameter of jax.vmap, please'
         ' make sure to run your jitted function inside the mesh context'
         ' manager.*SingleDeviceSharding'):
       jax.jit(jax.vmap(f, spmd_axis_name='x'))(arr)