Remove _manual_axes from NamedSharding since we can now track the manual axes on the mesh.

PiperOrigin-RevId: 748534841

Author: yashk2810

jax/BUILD

@@ -609,6 +609,7 @@ pytype_strict_library(
         ":dtypes",
         ":effects",
         ":layout",
+        ":mesh",
         ":op_shardings",
         ":partial_eval",
         ":partition_spec",

jax/_src/interpreters/mlir.py

@@ -49,6 +49,7 @@
 from jax._src.interpreters import xla
 from jax._src.layout import AutoLayout, DeviceLocalLayout
 from jax._src.partition_spec import PartitionSpec
+from jax._src.mesh import AxisType
 from jax._src.sharding import Sharding as JSharding
 from jax._src.sharding_impls import (AUTO, NamedSharding,
                                      modify_sdy_sharding_wrt_axis_types,
@@ -1017,18 +1018,29 @@ class LoweringResult(NamedTuple):
 
 
 def add_manual_axes(axis_ctx: sharding_impls.SPMDAxisContext, sharding, ndim):
-  mesh = axis_ctx.mesh
+  mesh = axis_ctx.mesh.abstract_mesh
+  sharding_mesh = sharding.mesh.abstract_mesh
   if (isinstance(sharding, sharding_impls.NamedSharding) and
-      sharding.mesh.shape == mesh.shape):
-    return sharding_impls.NamedSharding(
-        sharding.mesh, sharding.spec, memory_kind=sharding.memory_kind,
-        _manual_axes=axis_ctx.manual_axes)
+      sharding_mesh.shape == mesh.shape):
+    out_mesh, spec = sharding_mesh, sharding.spec
   else:
-    spec = sharding_impls.parse_flatten_op_sharding(
+    out_mesh, spec = mesh, sharding_impls.parse_flatten_op_sharding(
         sharding._to_xla_hlo_sharding(ndim), mesh)[0]
-    return sharding_impls.NamedSharding(
-      mesh, spec, memory_kind=sharding.memory_kind,
-      _manual_axes=axis_ctx.manual_axes)
+
+  out_mesh = out_mesh.update_axis_types(
+      {a: AxisType.Manual for a in axis_ctx.manual_axes})
+  out = sharding_impls.NamedSharding(out_mesh, spec,
+                                     memory_kind=sharding.memory_kind)
+  manual_axes = out.mesh.manual_axes
+  if any(p in manual_axes for s in out.spec
+         if s is not None and s is not PartitionSpec.UNCONSTRAINED
+         for p in (s if isinstance(s, tuple) else (s,))):
+    raise ValueError(
+        f'pspec {out.spec} contains a manual axes {manual_axes} of mesh'
+        f' which is not allowed. If you are using a'
+        ' with_sharding_constraint under a shard_map, only use the'
+        ' mesh axis in PartitionSpec which are not manual.')
+  return out
 
 
 def _to_physical_op_sharding(

jax/_src/named_sharding.py

@@ -112,20 +112,17 @@ class NamedSharding(JSharding.Sharding):
   mesh: mesh_lib.Mesh | mesh_lib.AbstractMesh
   spec: PartitionSpec
   _memory_kind: str | None
-  _manual_axes: frozenset[MeshAxisName]
   _logical_device_ids: tuple[int, ...] | None
 
   @use_cpp_method()
   def __init__(
       self, mesh: mesh_lib.Mesh | mesh_lib.AbstractMesh, spec: PartitionSpec, *,
-      memory_kind: str | None = None, _manual_axes=frozenset(),
-      _logical_device_ids=None):
+      memory_kind: str | None = None, _logical_device_ids=None):
     self.mesh = mesh
     self.spec = spec
     self._memory_kind = memory_kind
-    self._manual_axes = _manual_axes
     self._logical_device_ids = _logical_device_ids
-    check_pspec(self.mesh, self.spec, self._manual_axes)
+    check_pspec(self.mesh, self.spec)
 
   def __repr__(self):
     mem = '' if self.memory_kind is None else f', memory_kind={self.memory_kind}'
@@ -137,7 +134,6 @@ def __repr__(self):
   def __reduce__(self):
     return (type(self), (self.mesh, self.spec),
             {'memory_kind': self.memory_kind,
-             '_manual_axes': self._manual_axes,
              '_logical_device_ids': self._logical_device_ids})
 
   @property
@@ -147,8 +143,7 @@ def memory_kind(self) -> str | None:
   def __hash__(self):
     if not hasattr(self, '_hash'):
       self._hash = hash(
-          (self.mesh, self.memory_kind, self.spec, self._manual_axes,
-           self._logical_device_ids))
+          (self.mesh, self.memory_kind, self.spec, self._logical_device_ids))
     return self._hash
 
   def __eq__(self, other):
@@ -158,7 +153,6 @@ def __eq__(self, other):
       return True
     if (self.spec != other.spec
         or self.memory_kind != other.memory_kind
-        or self._manual_axes != other._manual_axes
         or self._logical_device_ids != other._logical_device_ids):
       return False
     return self.mesh is other.mesh or self.mesh == other.mesh
@@ -333,9 +327,7 @@ def named_sharding_to_xla_hlo_sharding(
   mesh_axis_pos = {name: i for i, name in enumerate(self.mesh.axis_names)}
 
   special_axes = {}
-  mesh_manual_axes = {n for n, t in self.mesh._name_to_type.items()
-                      if t == mesh_lib.AxisType.Manual}
-  manual_axes = self._manual_axes.union(mesh_manual_axes)
+  manual_axes = frozenset(self.mesh.manual_axes)
   if manual_axes:
     axis_names = self.mesh.axis_names
     for manual_axis in manual_axes:
@@ -420,7 +412,7 @@ def array_mapping_to_axis_resources(array_mapping: ArrayMapping):
 @cache(max_size=128, trace_context_in_key=False)
 def check_pspec(mesh, spec, _manual_axes=frozenset()):
   _check_unique_resources(spec, "NamedSharding spec", mesh)
-  _check_mesh_resource_axis(mesh, spec, _manual_axes)
+  _check_mesh_resource_axis(mesh, spec)
 
 class DuplicateSpecError(Exception):
   def __init__(self, message, mesh, pspec):
@@ -455,7 +447,7 @@ def _check_unique_resources(pspec: PartitionSpec, arg_name: str, mesh=None
         mesh=mesh, pspec=pspec)
 
 
-def _check_mesh_resource_axis(mesh, pspec, _manual_axes):
+def _check_mesh_resource_axis(mesh, pspec):
   for p in pspec:
     if p is PartitionSpec.UNCONSTRAINED or p is None:
       continue
@@ -465,10 +457,6 @@ def _check_mesh_resource_axis(mesh, pspec, _manual_axes):
         raise ValueError(
             f"Resource axis: {r} of {pspec} "
             f"is not found in mesh: {tuple(mesh.shape.keys())}.")
-      if r in _manual_axes:
-        raise ValueError(
-            f"Axis: {r} of {pspec} "
-            f"is also found in manual_axes: {_manual_axes}.") from None
     if not all(mesh._name_to_type[p[0]] == mesh._name_to_type[r] for r in p):
       raise ValueError(
           'AxisTypes should be the same in a tuple subset of PartitionSpec:'

jax/experimental/shard_map.py

@@ -785,13 +785,12 @@ def _shard_map_lowering(ctx, *in_nodes, jaxpr, mesh, in_names, out_names,
 
 def _make_scoped_manual_sharding(ctx, mesh, axes):
   axis_ctx = ctx.module_context.axis_context
+  mesh = mesh.abstract_mesh
   if isinstance(axis_ctx, sharding_impls.SPMDAxisContext):
-    manual_axes = axis_ctx.manual_axes
-  else:
-    manual_axes = frozenset({})
+    mesh = mesh.update_axis_types(
+        {a: AxisType.Manual for a in axis_ctx.manual_axes})
   return NamedSharding(
-      mesh, sharding_impls.array_mapping_to_axis_resources(axes),  # pytype: disable=wrong-arg-types
-      _manual_axes=manual_axes)
+      mesh, sharding_impls.array_mapping_to_axis_resources(axes))  # type: ignore
 
 def _xla_shard(ctx: mlir.LoweringRuleContext, mesh, auto, names,
                aval_in, aval_out, x):

jaxlib/xla/sharding.cc

@@ -167,8 +167,6 @@ bool ShardingEqual(nb::handle a, nb::handle b) {
            a_named_sharding->spec().equal(b_named_sharding->spec()) &&
            a_named_sharding->memory_kind().equal(
                b_named_sharding->memory_kind()) &&
-           a_named_sharding->manual_axes().equal(
-               b_named_sharding->manual_axes()) &&
            a_named_sharding->logical_device_ids().equal(
                b_named_sharding->logical_device_ids());
   }
@@ -204,15 +202,14 @@ static const std::array<absl::string_view, 3> valid_memory_kinds = {
 };
 
 NamedSharding::NamedSharding(nb::object mesh, nb::object spec,
-                             nb::object memory_kind, nb::object manual_axes,
+                             nb::object memory_kind,
                              nb::object logical_device_ids)
     : Sharding(/*num_devices=*/[&mesh]() {
         return nb::cast<int>(mesh.attr("size"));
       }()),
       mesh_(std::move(mesh)),
       spec_(std::move(spec)),
       memory_kind_(std::move(memory_kind)),
-      manual_axes_(std::move(manual_axes)),
       logical_device_ids_(std::move(logical_device_ids)) {
   if (spec_.is_none()) {
     throw nb::type_error(
@@ -261,7 +258,7 @@ NamedSharding::NamedSharding(nb::object mesh, nb::object spec,
     }
     return output;
   }();
-  (*check_pspec)(mesh_, spec_, manual_axes_);
+  (*check_pspec)(mesh_, spec_);
 }
 
 /*static*/ PyObject* NamedSharding::type_ = nullptr;
@@ -352,16 +349,13 @@ void RegisterSharding(nb::module_& m) {
   nb::class_<Sharding>(m, "Sharding").def(nb::init<>());
 
   nb::class_<NamedSharding, Sharding>(m, "NamedSharding", nb::dynamic_attr())
-      .def(nb::init<nb::object, nb::object, nb::object, nb::object,
-                    nb::object>(),
+      .def(nb::init<nb::object, nb::object, nb::object, nb::object>(),
            nb::arg("mesh"), nb::arg("spec").none(),
            nb::arg("memory_kind").none() = nb::none(),
-           nb::arg("_manual_axes") = nb::steal(PyFrozenSet_New(nullptr)),
            nb::arg("_logical_device_ids").none() = nb::none())
       .def_prop_ro("mesh", &NamedSharding::mesh)
       .def_prop_ro("spec", &NamedSharding::spec)
       .def_prop_ro("_memory_kind", &NamedSharding::memory_kind)
-      .def_prop_ro("_manual_axes", &NamedSharding::manual_axes)
       .def_prop_ro("_logical_device_ids", &NamedSharding::logical_device_ids)
       .def_prop_ro("_internal_device_list", [](const NamedSharding& s) {
         return xla::ValueOrThrow(s.internal_device_list());

jaxlib/xla/sharding.h

@@ -75,13 +75,12 @@ bool ShardingEqual(nanobind::handle a, nanobind::handle b);
 class NamedSharding : public Sharding {
  public:
   NamedSharding(nanobind::object mesh, nanobind::object spec,
-                nanobind::object memory_kind, nanobind::object manual_axes,
+                nanobind::object memory_kind,
                 nanobind::object logical_device_ids);
 
   const nanobind::object& mesh() const { return mesh_; }
   const nanobind::object& spec() const { return spec_; }
   const nanobind::object& memory_kind() const { return memory_kind_; }
-  const nanobind::object& manual_axes() const { return manual_axes_; }
   const nanobind::object& logical_device_ids() const {
     return logical_device_ids_;
   }
@@ -102,7 +101,6 @@ class NamedSharding : public Sharding {
   nanobind::object mesh_;
   nanobind::object spec_;
   nanobind::object memory_kind_;
-  nanobind::object manual_axes_;
   nanobind::object logical_device_ids_;
   std::optional<xla::nb_class_ptr<PyDeviceList>> internal_device_list_;
   static PyObject* type_;

jaxlib/xla/xla_extension/__init__.pyi

@@ -926,14 +926,12 @@ class NamedSharding(Sharding):
       spec: Any,
       *,
       memory_kind: Optional[str] = None,
-      _manual_axes: frozenset[Any] = frozenset(),
       _logical_device_ids: tuple[int, ...] | None = None,
   ): ...
   mesh: Any
   spec: Any
   _memory_kind: Optional[str]
   _internal_device_list: DeviceList
-  _manual_axes: frozenset[Any]
   _logical_device_ids: tuple[int, ...] | None
 
 class SingleDeviceSharding(Sharding):

tests/shard_map_test.py

@@ -2234,17 +2234,12 @@ def g(x):
       return x * x
 
     def h(x):
-      return shard_map(g, mesh,
-                    in_specs=P(None, 'j'),
-                    out_specs=P(None, 'j'))(x)
+      return shard_map(g, mesh, in_specs=P(None, 'j'), out_specs=P(None, 'j'))(x)
 
     @jax.jit
     def f(x):
-      return shard_map(h, mesh,
-                    in_specs=P('i', None),
-                    out_specs=P('i', None),
-                    check_rep=False,
-                    auto=frozenset({'j'}))(x)
+      return shard_map(h, mesh, in_specs=P('i', None), out_specs=P('i', None),
+                       check_rep=False, auto=frozenset({'j'}))(x)
 
     v = jnp.arange(32.).reshape(4, 8)
     v = jax.device_put(v, jax.sharding.NamedSharding(mesh, P('i', 'j')))