spmd_types 0.2.2__tar.gz → 0.2.3__tar.gz

{spmd_types-0.2.2/spmd_types.egg-info → spmd_types-0.2.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: spmd_types
-Version: 0.2.2
+Version: 0.2.3
 Summary: A type system for distributed (SPMD) tensor computations in PyTorch
 Author: Meta Platforms, Inc.
 License: BSD 3-Clause License
@@ -47,7 +47,6 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: torch>=2.10.0
 Dynamic: license-file
 
 # `spmd_types`

{spmd_types-0.2.2 → spmd_types-0.2.3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "spmd_types"
-version = "0.2.2"
+version = "0.2.3"
 description = "A type system for distributed (SPMD) tensor computations in PyTorch"
 readme = "README.md"
 license = {file = "LICENSE"}
@@ -24,7 +24,6 @@ classifiers = [
     "Topic :: Scientific/Engineering :: Artificial Intelligence",
 ]
 dependencies = [
-    "torch>=2.10.0",
 ]
 
 [project.urls]

{spmd_types-0.2.2 → spmd_types-0.2.3}/spmd_types/__init__.py

@@ -41,7 +41,6 @@ from spmd_types._state import (  # noqa: F401
     current_mesh_all_names,
     current_mesh_names,
     is_type_checking,
-    no_typecheck,
 )
 from spmd_types._traceback import traceback_filtering  # noqa: F401
 from spmd_types._type_attr import (  # noqa: F401
@@ -56,10 +55,13 @@ from spmd_types.runtime import (  # noqa: F401
     assert_local_type_like,
     assert_type,
     assert_type_like,
+    get_partition_spec,
+    has_local_type,
     Infer,
     local,
     local_map,
     mutate_type,
+    no_typecheck,
     register_autograd_function,
     register_decomposition,
     register_local_autograd_function,

{spmd_types-0.2.2 → spmd_types-0.2.3}/spmd_types/_checker/__init__.py

@@ -68,6 +68,7 @@ from spmd_types.runtime import (  # noqa: F401
     get_partition_spec,
     has_local_type,
     mutate_type,
+    no_typecheck,
     register_autograd_function,
     register_decomposition,
     register_local_autograd_function,
@@ -2531,6 +2532,8 @@ def _get_autograd_function_class(func) -> type | None:
 def _validate_partition_spec_for_global_spmd(
     local_type: LocalSpmdType,
     spec: PartitionSpec | None,
+    *,
+    is_global_axis: Callable[[MeshAxis], bool],
 ) -> None:
     """Validate PartitionSpec <-> local type consistency.
 
@@ -2543,10 +2546,12 @@ def _validate_partition_spec_for_global_spmd(
     Args:
         local_type: The tensor's local SPMD type dict.
         spec: The tensor's PartitionSpec (None if absent).
+        is_global_axis: Returns whether an axis requires global PartitionSpec
+            validation in the current mesh context.
     """
     spec_axes = spec.axes_with_partition_spec() if spec is not None else set()
     for axis, typ in local_type.items():
-        if typ is V and axis not in spec_axes:
+        if is_global_axis(axis) and typ is V and axis not in spec_axes:
             ax = format_axis(axis)
             raise SpmdTypeError(
                 f"Tensor has Varying type on axis {ax} but no PartitionSpec entry. "
@@ -2555,7 +2560,7 @@ def _validate_partition_spec_for_global_spmd(
             )
     for axis in spec_axes:
         typ = local_type.get(axis)
-        if typ is not V:
+        if is_global_axis(axis) and typ is not V:
             ax = format_axis(axis)
             raise SpmdTypeError(
                 f"Tensor has PartitionSpec on axis {ax} but local type is "
@@ -2623,6 +2628,23 @@ class _SpmdTypeMode(torch.overrides.TorchFunctionMode):
 
         self._disabled = False
 
+    def _is_global_axis(self, axis: MeshAxis) -> bool:
+        if self._local:
+            return False
+        local_axes = _state._current_local_axes()
+        if not local_axes:
+            return True
+        mesh = _state.current_mesh()
+        if mesh is None or axis not in mesh:
+            raise SpmdTypeError(
+                f"Axis {format_axis(axis)} is not in the current mesh, so its "
+                "local/global typechecking policy is undefined"
+            )
+        return axis not in local_axes
+
+    def _global_axes_for(self, axes: Iterable[MeshAxis]) -> set[MeshAxis]:
+        return {axis for axis in axes if self._is_global_axis(axis)}
+
     def __torch_function__(self, func, types, args=(), kwargs=None):  # noqa: C901
         kwargs = kwargs or {}
         # Paused via no_typecheck(): run without type checking.
@@ -2679,7 +2701,11 @@ class _SpmdTypeMode(torch.overrides.TorchFunctionMode):
             # Global SPMD validation.
             if not self._local:
                 for local_type, spec in zip(info.tensor_types, info.partition_specs):
-                    _validate_partition_spec_for_global_spmd(local_type, spec)
+                    _validate_partition_spec_for_global_spmd(
+                        local_type,
+                        spec,
+                        is_global_axis=self._is_global_axis,
+                    )
 
             # Typecheck-registered autograd function: dispatch to typecheck_forward
             # INSTEAD of executing func.  The mode is popped off the
@@ -2744,6 +2770,29 @@ class _SpmdTypeMode(torch.overrides.TorchFunctionMode):
                 # Resolve axis to a list of matching axes on the tensor.
                 # Direct match: [axis]. Flattened decomposition: [sub1, sub2, ...].
                 resolved_axes = _resolve_collective_axes(axis, local_type)
+                collective_axes = resolved_axes or [axis]
+                global_collective_axes = [
+                    ax for ax in collective_axes if self._is_global_axis(ax)
+                ]
+                if global_collective_axes and len(global_collective_axes) != len(
+                    collective_axes
+                ):
+                    global_axis_set = set(global_collective_axes)
+                    local_collective_axes = [
+                        ax for ax in collective_axes if ax not in global_axis_set
+                    ]
+                    local_names = ", ".join(
+                        format_axis(ax) for ax in local_collective_axes
+                    )
+                    global_names = ", ".join(
+                        format_axis(ax) for ax in global_collective_axes
+                    )
+                    raise SpmdTypeError(
+                        f"Collective axis {format_axis(axis)} spans mesh axes with "
+                        f"mixed typechecking policies: local axes ({local_names}) "
+                        f"and global axes ({global_names}). Pass a collective axis "
+                        f"whose component axes are all local or all global."
+                    )
 
                 if resolved_axes is not None:
                     # Validate input types on all resolved axes
@@ -2774,7 +2823,7 @@ class _SpmdTypeMode(torch.overrides.TorchFunctionMode):
                         output_type[ax] = local_dst
                 _set_local_type(result, output_type)
 
-                if _state._is_global():
+                if global_collective_axes:
                     input_spec = get_partition_spec(x)
                     # Find S(dim) for this axis in the input spec, validating
                     # it is innermost in its multi-axis group.
@@ -2916,15 +2965,16 @@ class _SpmdTypeMode(torch.overrides.TorchFunctionMode):
                     )
 
                 # Stage 3: global shard propagation.
-                if self._local:
-                    global_shard_axes, shard_edges = [], {}
-                else:
-                    all_axes: set[DeviceMeshAxis] = set()
-                    for typ in input_types_list:
-                        all_axes.update(typ.keys())
+                all_axes: set[MeshAxis] = set()
+                for typ in input_types_list:
+                    all_axes.update(typ.keys())
+                global_axes = self._global_axes_for(all_axes)
+                if global_axes:
                     global_shard_axes, shard_edges = _collect_shard_axes(
-                        partition_specs, all_axes
+                        partition_specs, global_axes
                     )
+                else:
+                    global_shard_axes, shard_edges = [], {}
 
                 # Validate mutation safety for in-place/out operations.
                 mutated = _get_mutated_tensors(func, args, kwargs, result)
@@ -3086,9 +3136,6 @@ def typecheck(
             _remove_autograd_apply_patch()
 
 
-from spmd_types._state import no_typecheck  # noqa: F401
-
-
 class _SpmdTypeBackwardCompatibleMode:
     """Backwards-compatible wrapper around ``typecheck()``.
 

{spmd_types-0.2.2 → spmd_types-0.2.3}/spmd_types/_mesh.py

@@ -13,8 +13,8 @@ from contextlib import contextmanager
 from typing import TYPE_CHECKING
 
 from spmd_types._mesh_axis import MeshAxis
-from spmd_types._state import _axes_to_pgs, _pop_mesh, _push_mesh
-from spmd_types.types import DeviceMeshAxis, normalize_mesh
+from spmd_types._state import _axes_to_pgs, _current_mesh_entry, _pop_mesh, _push_mesh
+from spmd_types.types import DeviceMeshAxis, format_axis, normalize_axis, normalize_mesh
 from torch.distributed.device_mesh import DeviceMesh
 
 if TYPE_CHECKING:
@@ -24,8 +24,14 @@ if TYPE_CHECKING:
 @contextmanager
 def set_current_mesh(
     axes: (
-        frozenset[MeshAxis] | DeviceMesh | Sequence[ProcessGroup] | dict[str, MeshAxis]
-    ),
+        frozenset[MeshAxis]
+        | DeviceMesh
+        | Sequence[ProcessGroup]
+        | dict[str, MeshAxis]
+        | None
+    ) = None,
+    *,
+    local_axes: tuple[DeviceMeshAxis, ...] = (),
 ):
     """Context manager that pushes a mesh onto the stack.
 
@@ -37,19 +43,99 @@ def set_current_mesh(
             - A frozenset of orthogonal MeshAxis objects (no string lookup).
             - A sequence of ProcessGroup objects, each converted via
               ``MeshAxis.of()`` (no string lookup).
+            When omitted, reuses the current mesh.
+        local_axes: An outer prefix of mesh axes that retains local SPMD
+            semantics during global type checking. Each coordinate on these
+            axes selects an independent local tensor; the remaining inner axes
+            describe the global sharding of that tensor. For example, with
+            ``dp`` local and ``tp`` global, each ``dp`` coordinate has a
+            separate tensor globally sharded across ``tp``.
 
     Singleton (size-1) axes are dropped from the active axis set via
     ``normalize_mesh`` but remain in the name lookup table
     (``current_mesh_all_names``).
     """
-    resolved, names = _resolve_axes(axes)
-    _push_mesh(resolved, names, _axes_to_pgs(axes))
+    if axes is None:
+        entry = _current_mesh_entry()
+        if entry is None:
+            raise RuntimeError(
+                "set_current_mesh() without axes requires a current mesh"
+            )
+        resolved, names, pgs = entry.axes, entry.all_names, entry.pgs
+    else:
+        resolved, names = _resolve_axes(axes)
+        pgs = _axes_to_pgs(axes)
+    resolved_local_axes = _resolve_local_axes(local_axes, resolved, names)
+    _push_mesh(
+        resolved,
+        names,
+        pgs,
+        local_axes=resolved_local_axes,
+    )
     try:
         yield
     finally:
         _pop_mesh()
 
 
+def _resolve_local_axes(
+    local_axes: tuple[DeviceMeshAxis, ...],
+    mesh: frozenset[MeshAxis],
+    names: dict[str, MeshAxis],
+) -> frozenset[MeshAxis]:
+    """Resolve ``local_axes`` to members of ``mesh``.
+
+    Rejects unknown or duplicate axes and requires the local axes to be the
+    outermost axes and the global axes to be the innermost axes. For a
+    ``(dp, tp)`` mesh, ``("dp",)`` is valid but ``("tp",)`` is not.
+
+    Returns the resolved local axes as a ``frozenset[MeshAxis]``.
+    """
+    unknown = tuple(
+        axis for axis in local_axes if isinstance(axis, str) and axis not in names
+    )
+    if unknown:
+        raise ValueError(
+            f"local_axes contains names not present in this mesh: {unknown!r}; "
+            f"available names: {tuple(sorted(names))!r}"
+        )
+    resolved_list = tuple(
+        names[axis] if isinstance(axis, str) else normalize_axis(axis)
+        for axis in local_axes
+    )
+    if len(set(resolved_list)) != len(resolved_list):
+        duplicates = {axis for axis in resolved_list if resolved_list.count(axis) > 1}
+        duplicate_names = ", ".join(format_axis(axis) for axis in duplicates)
+        raise ValueError(
+            f"local_axes must not contain duplicate axes; got {local_axes!r}; "
+            f"duplicate axes: ({duplicate_names})"
+        )
+    resolved = frozenset(resolved_list)
+    if not resolved <= mesh:
+        outside = resolved - mesh
+        outside_names = ", ".join(format_axis(axis) for axis in outside)
+        mesh_names = ", ".join(format_axis(axis) for axis in sorted(mesh, key=repr))
+        raise ValueError(
+            f"local_axes contains axes outside this mesh: ({outside_names}); "
+            f"current mesh axes: ({mesh_names})"
+        )
+
+    ordered = sorted(
+        mesh,
+        key=lambda axis: max(stride for _, stride in axis.layout.sizes_and_strides),
+        reverse=True,
+    )
+    prefix = set(ordered[: len(resolved)])
+    if set(resolved) != prefix:
+        local_names = ", ".join(format_axis(axis) for axis in resolved)
+        mesh_names = ", ".join(format_axis(axis) for axis in ordered)
+        raise ValueError(
+            "local_axes must come before all global axes in outer-to-inner mesh "
+            f"order; got local_axes=({local_names}) for mesh=({mesh_names})"
+        )
+    return resolved
+
+
 def _pg_for_axis(axis: DeviceMeshAxis) -> ProcessGroup:
     """Resolve a mesh axis to its ``ProcessGroup`` using the current mesh.
 
@@ -57,14 +143,13 @@ def _pg_for_axis(axis: DeviceMeshAxis) -> ProcessGroup:
     ``MeshAxis``. Raises ``RuntimeError`` if no current mesh carries process
     groups, or if the axis is not part of it.
     """
-    from spmd_types._state import _tls, current_mesh_all_names
     from torch.distributed import ProcessGroup
 
     if isinstance(axis, ProcessGroup):
         return axis
 
-    stack = getattr(_tls, "mesh_stack", None)
-    pgs = stack[-1].pgs if stack else None
+    entry = _current_mesh_entry()
+    pgs = entry.pgs if entry is not None else None
     if not pgs:
         raise RuntimeError(
             "Passing an axis name to a redistribution or collective "
@@ -73,7 +158,7 @@ def _pg_for_axis(axis: DeviceMeshAxis) -> ProcessGroup:
         )
 
     if isinstance(axis, str):
-        resolved = (current_mesh_all_names() or {}).get(axis)
+        resolved = entry.all_names.get(axis) if entry is not None else None
     else:
         resolved = axis
     if resolved is None or resolved not in pgs:

{spmd_types-0.2.2 → spmd_types-0.2.3}/spmd_types/_state.py

@@ -30,6 +30,7 @@ if TYPE_CHECKING:
 
 class MeshEntry(NamedTuple):
     axes: frozenset[MeshAxis]
+    local_axes: frozenset[MeshAxis]
     names: dict[str, MeshAxis]
     all_names: dict[str, MeshAxis]
     pgs: dict[MeshAxis, ProcessGroup]
@@ -45,7 +46,7 @@ def is_type_checking() -> bool:
 
 
 @contextmanager
-def no_typecheck():
+def _no_typecheck_context():
     """Temporarily disable type checking on this thread (like ``no_grad``)."""
     mode = _current_mode()
     if mode is not None:
@@ -75,12 +76,16 @@ def is_strict() -> bool:
     return mode is not None and mode._strict
 
 
+def _current_mesh_entry() -> MeshEntry | None:
+    """Return the active mesh-stack entry, or None if no mesh is set."""
+    stack: list[MeshEntry] = getattr(_tls, "mesh_stack", [])
+    return stack[-1] if stack else None
+
+
 def current_mesh() -> frozenset[MeshAxis] | None:
     """Return the current mesh axes, or None if no mesh is set."""
-    stack = getattr(_tls, "mesh_stack", None)
-    if stack:
-        return stack[-1].axes
-    return None
+    entry = _current_mesh_entry()
+    return entry.axes if entry is not None else None
 
 
 def current_mesh_names() -> dict[str, MeshAxis] | None:
@@ -89,18 +94,14 @@ def current_mesh_names() -> dict[str, MeshAxis] | None:
     Only includes non-singleton axes, consistent with ``current_mesh()``.
     Use ``current_mesh_all_names()`` to include singleton axes.
     """
-    stack = getattr(_tls, "mesh_stack", None)
-    if stack:
-        return stack[-1].names
-    return None
+    entry = _current_mesh_entry()
+    return entry.names if entry is not None else None
 
 
 def current_mesh_all_names() -> dict[str, MeshAxis] | None:
     """Return the full name-to-axis mapping including singleton axes, or None."""
-    stack = getattr(_tls, "mesh_stack", None)
-    if stack:
-        return stack[-1].all_names
-    return None
+    entry = _current_mesh_entry()
+    return entry.all_names if entry is not None else None
 
 
 def _axes_to_pgs(
@@ -133,6 +134,7 @@ def _push_mesh(
     | dict[str, MeshAxis],
     names: dict[str, MeshAxis] | None = None,
     pgs: dict[MeshAxis, ProcessGroup] | None = None,
+    local_axes: frozenset[MeshAxis] = frozenset(),
 ) -> None:
     if isinstance(mesh, frozenset):
         resolved, resolved_names = mesh, names if names is not None else {}
@@ -149,7 +151,9 @@ def _push_mesh(
     filtered_names = {k: v for k, v in resolved_names.items() if v in resolved}
     if not hasattr(_tls, "mesh_stack"):
         _tls.mesh_stack = []
-    _tls.mesh_stack.append(MeshEntry(resolved, filtered_names, all_names, resolved_pgs))
+    _tls.mesh_stack.append(
+        MeshEntry(resolved, local_axes, filtered_names, all_names, resolved_pgs)
+    )
 
 
 def _pop_mesh() -> MeshEntry:
@@ -166,6 +170,12 @@ def _clear_mesh_stack() -> None:
     _tls.mesh_stack = []
 
 
+def _current_local_axes() -> frozenset[MeshAxis]:
+    """Return the local axes for the active mesh, or an empty set."""
+    entry = _current_mesh_entry()
+    return entry.local_axes if entry is not None else frozenset()
+
+
 def _is_global() -> bool:
     """Return True if global SPMD mode is active (local=False)."""
     mode = _current_mode()

{spmd_types-0.2.2 → spmd_types-0.2.3}/spmd_types/checker.py

@@ -9,5 +9,5 @@ from spmd_types._checker import (  # noqa: F401
     _SpmdTypeBackwardCompatibleMode as SpmdTypeMode,
     typecheck,
 )
-from spmd_types._state import is_type_checking, no_typecheck  # noqa: F401
-from spmd_types.runtime import local_map  # noqa: F401
+from spmd_types._state import is_type_checking  # noqa: F401
+from spmd_types.runtime import local_map, no_typecheck  # noqa: F401

{spmd_types-0.2.2 → spmd_types-0.2.3}/spmd_types/runtime.py

@@ -27,14 +27,14 @@ from __future__ import annotations
 import builtins
 import logging
 import os
-from contextlib import contextmanager
-from typing import Any, TypeAlias
+from contextlib import AbstractContextManager, contextmanager
+from typing import Any, Callable, overload, TypeAlias
 
 import torch
 from spmd_types._frame import _get_user_frame
 from spmd_types._mesh_axis import MeshAxis
 from spmd_types._scalar_sentinel import _Scalar
-from spmd_types._state import current_mesh, is_type_checking
+from spmd_types._state import _no_typecheck_context, current_mesh, is_type_checking
 from spmd_types._traceback import api_boundary
 from spmd_types._type_attr import (
     _LOCAL_TYPE_ATTR,
@@ -589,20 +589,49 @@ def assert_local_type_like(
 # =============================================================================
 
 
-@api_boundary
-def mutate_type(
-    tensor: torch.Tensor,
+def _mutate_type_single_axis(
+    local_type: LocalSpmdType,
     axis: DeviceMeshAxis,
     *,
     src: PerMeshAxisSpmdType,
     dst: PerMeshAxisSpmdType,
+) -> LocalSpmdType:
+    axis = normalize_axis(axis)
+    if axis.size() == 1:
+        return local_type
+    if axis not in local_type:
+        raise SpmdTypeError(
+            f"mutate_type: axis {format_axis(axis)} not found in tensor's "
+            f"SPMD type. Tensor has axes: {set(local_type.keys())}"
+        )
+
+    src_local = to_local_type(src)
+    current = local_type[axis]
+    if current is not src_local:
+        raise SpmdTypeError(
+            f"mutate_type: expected current type {src_local} on axis "
+            f"{format_axis(axis)}, got {current}"
+        )
+
+    new_type = dict(local_type)
+    new_type[axis] = to_local_type(dst)
+    return new_type
+
+
+@api_boundary
+def mutate_type(
+    tensor: torch.Tensor,
+    axis: DeviceMeshAxis | None = None,
+    *,
+    src: PerMeshAxisSpmdType | dict[DeviceMeshAxis, PerMeshAxisSpmdType],
+    dst: PerMeshAxisSpmdType | dict[DeviceMeshAxis, PerMeshAxisSpmdType],
 ) -> torch.Tensor:
-    """Change the SPMD type of a single mesh axis on an already-annotated tensor.
+    """Change the SPMD type of mesh axes on an already-annotated tensor.
 
     Unlike ``assert_type``, this function *overwrites* the existing type on
-    ``axis``.  The caller must specify the expected current type (``src``) to
-    prevent silent corruption; a ``SpmdTypeError`` is raised if the tensor's
-    current type on ``axis`` does not match ``src``.
+    the updated axes. The caller must specify the expected current type
+    (``src``) to prevent silent corruption; a ``SpmdTypeError`` is raised if
+    an updated axis's current type does not match its expected ``src``.
 
     This is intended for internal use in low-level parallelism primitives
     where a buffer legitimately changes its distribution semantics in place
@@ -613,45 +642,88 @@ def mutate_type(
     *views* into a buffer after an in-place collective has changed the
     buffer's semantics.
 
+    There are four supported call shapes:
+
+    - ``mutate_type(tensor, axis, src=R, dst=I)`` updates a single axis.
+    - ``mutate_type(tensor, src=R, dst=I)`` updates all axes annotated on the
+      tensor, assuming all axes share ``src`` type.
+    - ``mutate_type(tensor, src=R, dst={axis0: I, axis1: P})`` updates every
+      axis in ``dst``, after checking all update axes have ``src`` type.
+      Axes not listed in ``dst`` are unaffected.
+    - ``mutate_type(tensor, src={axis0: R, axis1: I}, dst=P)`` is similarly supported.
+    - ``mutate_type(tensor, src={axis0: R, axis1: R}, dst={axis0: I, axis1: I})``
+      updates each axis correspondingly. ``src`` and ``dst`` keys must match,
+      and other axes are unaffected.
+
+    ``mutate_type`` errors if any specified ``src`` type doesn't match the tensor's
+    annotation.
+
     Args:
         tensor: The tensor whose type to mutate.  Must already have a local
             type annotation.
-        axis: The mesh axis (MeshAxis or ProcessGroup) to modify.
-        src: The expected current type on ``axis``.  Raises if it does not
-            match.  Accepts R, I, V, P, or S(i) (S(i) is compared as V).
-        dst: The new type to set on ``axis``.  Accepts R, I, V, P, or S(i)
-            (S(i) is stored as V).
+        axis: The mesh axis (MeshAxis or ProcessGroup) to modify for the
+            single-axis form. Omit when ``src`` or ``dst`` is a dict mapping
+            axes to types.
+        src: The expected current type, or a dict of per-axis expected current
+            types. Accepts R, I, V, P, or S(i); S(i) is compared as V.
+        dst: The new type, or a dict of per-axis destination types. Accepts R,
+            I, V, P, or S(i); S(i) is stored as V.
 
     Returns:
         The tensor for chaining.
 
     Raises:
-        SpmdTypeError: If the tensor has no type, the axis is missing,
-            or the current type does not match ``src``.
+        SpmdTypeError: If the tensor has no type, if a non-singleton updated
+            axis is missing from the tensor's local type, if a current type
+            does not match ``src``, or if per-axis ``src`` and ``dst`` dict keys
+            differ.
     """
-    axis = normalize_axis(axis)
-    if axis.size() == 1:
-        return tensor  # singleton axes are not tracked
     local_type = get_local_type(tensor)
-    if axis not in local_type:
-        raise SpmdTypeError(
-            f"mutate_type: axis {format_axis(axis)} not found in tensor's "
-            f"SPMD type. Tensor has axes: {set(local_type.keys())}"
-        )
-
-    # Normalize S(i) -> V for comparison and storage
-    src_local = to_local_type(src)
-    dst_local = to_local_type(dst)
-
-    current = local_type[axis]
-    if current is not src_local:
-        raise SpmdTypeError(
-            f"mutate_type: expected current type {src_local} on axis "
-            f"{format_axis(axis)}, got {current}"
+    if isinstance(dst, dict) and isinstance(src, dict):
+        if axis is not None:
+            raise SpmdTypeError(
+                "mutate_type() expected axis to be omitted when src or dst is a dict."
+            )
+        if set(dst) != set(src):
+            axes = ", ".join(format_axis(axis) for axis in set(dst) ^ set(src))
+            raise SpmdTypeError(
+                f"mutate_type() expected src and dst axes to match: {axes}"
+            )
+        updates = [
+            (update_axis, src[update_axis], dst_type)
+            for update_axis, dst_type in dst.items()
+        ]
+    elif isinstance(dst, dict):
+        if axis is not None:
+            raise SpmdTypeError(
+                "mutate_type() expected axis to be omitted when src or dst is a dict."
+            )
+        updates = [
+            (update_axis, src, dst_type) for update_axis, dst_type in dst.items()
+        ]
+    elif isinstance(src, dict):
+        if axis is not None:
+            raise SpmdTypeError(
+                "mutate_type() expected axis to be omitted when src or dst is a dict."
+            )
+        updates = [
+            (update_axis, src_type, dst) for update_axis, src_type in src.items()
+        ]
+    else:
+        updates = (
+            [(axis, src, dst)]
+            if axis is not None
+            else [(update_axis, src, dst) for update_axis in local_type.keys()]
         )
 
     new_type = dict(local_type)
-    new_type[axis] = dst_local
+    for update_axis, src_type, dst_type in updates:
+        new_type = _mutate_type_single_axis(
+            new_type,
+            update_axis,
+            src=src_type,
+            dst=dst_type,
+        )
     return _set_local_type(tensor, new_type)
 
 
@@ -908,7 +980,32 @@ def local():
         yield
 
 
-def local_map(  # noqa: C901
+@overload
+def no_typecheck() -> AbstractContextManager[None]: ...
+
+
+@overload
+def no_typecheck(
+    *, out_types: Any, in_types: Any = ...
+) -> Callable[[Callable[..., Any]], Callable[..., Any]]: ...
+
+
+def no_typecheck(**kwargs: Any):
+    """Disable checking in a context or across an explicitly typed function."""
+    if not kwargs:
+        return _no_typecheck_context()
+    if "out_types" not in kwargs or kwargs.keys() - {"in_types", "out_types"}:
+        raise TypeError("out_types is required for typed no_typecheck")
+
+    return _typecheck_boundary(
+        name="no_typecheck",
+        body_context=_no_typecheck_context,
+        in_types=kwargs.get("in_types", Infer),
+        out_types=kwargs["out_types"],
+    )
+
+
+def local_map(
     *,
     in_types: Any = Infer,
     out_types: Any,
@@ -980,12 +1077,23 @@ def local_map(  # noqa: C901
         def fn(x, weights):
             return x @ weights[0], x @ weights[1]
     """
+    return _typecheck_boundary(
+        name="local_map",
+        body_context=local,
+        in_types=in_types,
+        out_types=out_types,
+    )
+
+
+def _typecheck_boundary(  # noqa: C901
+    *, name: str, body_context: Any, in_types: Any, out_types: Any
+):
     import functools
 
     from torch.utils._pytree import tree_map
 
     def decorator(fn):
-        prefix = f"local_map[{fn.__qualname__}]"
+        prefix = f"{name}[{fn.__qualname__}]"
 
         def _check_arg(arg: object, spec: _LocalMapSpecLeaf, kind: str) -> None:
             if spec is None:
@@ -1087,7 +1195,7 @@ def local_map(  # noqa: C901
             if in_types is not Infer:
                 _walk_boundary(args, in_types, "input")
 
-            with local():
+            with body_context():
                 result = fn(*args, **kwargs)
 
             _walk_boundary(result, out_types, "output")

{spmd_types-0.2.2 → spmd_types-0.2.3/spmd_types.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: spmd_types
-Version: 0.2.2
+Version: 0.2.3
 Summary: A type system for distributed (SPMD) tensor computations in PyTorch
 Author: Meta Platforms, Inc.
 License: BSD 3-Clause License
@@ -47,7 +47,6 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: torch>=2.10.0
 Dynamic: license-file
 
 # `spmd_types`

{spmd_types-0.2.2 → spmd_types-0.2.3}/spmd_types.egg-info/SOURCES.txt

@@ -29,6 +29,5 @@ spmd_types/types.py
 spmd_types.egg-info/PKG-INFO
 spmd_types.egg-info/SOURCES.txt
 spmd_types.egg-info/dependency_links.txt
-spmd_types.egg-info/requires.txt
 spmd_types.egg-info/top_level.txt
 spmd_types/_checker/__init__.py

spmd_types-0.2.2/spmd_types.egg-info/requires.txt

@@ -1 +0,0 @@
-torch>=2.10.0