spmd_types 0.2.2__tar.gz → 0.2.4__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: spmd_types
-Version: 0.2.2
+Version: 0.2.4
 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.4}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "spmd_types"
-version = "0.2.2"
+version = "0.2.4"
 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.4}/spmd_types/__init__.py

@@ -7,9 +7,6 @@
 # spmd_types package
 from __future__ import annotations
 
-from spmd_types._backward_hooks import (  # noqa: F401
-    register_local_backward_hook,
-)
 from spmd_types._collectives import (  # noqa: F401
     all_gather,
     all_reduce,
@@ -30,6 +27,10 @@ from spmd_types._local import (  # noqa: F401
     reinterpret,
     shard,
 )
+from spmd_types._local_registration import (  # noqa: F401
+    register_local_autograd_function,
+    register_local_backward_hook,
+)
 from spmd_types._mesh import set_current_mesh  # noqa: F401
 from spmd_types._mesh_axis import MeshAxis  # noqa: F401
 
@@ -41,7 +42,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,13 +56,15 @@ 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,
     trace,
 )
 

{spmd_types-0.2.2 → spmd_types-0.2.4}/spmd_types/_backward_hooks.py

@@ -18,10 +18,12 @@ motivates the exact shape.
 
 from __future__ import annotations
 
-from collections.abc import Callable
-
 import torch
 import torch.utils.hooks as _torch_hooks
+from spmd_types._local_registration import (  # noqa: F401
+    _LOCAL_BACKWARD_HOOKS,
+    register_local_backward_hook,
+)
 from spmd_types._state import is_type_checking
 from spmd_types.runtime import (
     _set_partition_spec,
@@ -31,20 +33,6 @@ from spmd_types.runtime import (
 from spmd_types.types import SpmdTypeError
 from torch.nn.modules._functions import BackwardHookFunction
 
-_LOCAL_BACKWARD_HOOKS: set[Callable] = set()
-
-
-def register_local_backward_hook(fn: Callable) -> Callable:
-    """Declare that ``fn`` does not alter SPMD types when it runs in backward.
-
-    Analogous to :func:`register_local_autograd_function`: the hook's backward
-    is treated as local (no collectives, no type-changing effects on grads).
-    Side-agnostic: covers both ``register_full_backward_hook`` (post) and
-    ``register_full_backward_pre_hook`` (pre) use.  Usable as a decorator.
-    """
-    _LOCAL_BACKWARD_HOOKS.add(fn)
-    return fn
-
 
 def _validate(hooks):
     if not is_type_checking():

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

@@ -36,6 +36,10 @@ from spmd_types._collectives import (
     reduce_scatter,
 )
 from spmd_types._local import convert, invariant_to_replicate, reinterpret
+from spmd_types._local_registration import (  # noqa: F401
+    _LOCAL_AUTOGRAD_FUNCTIONS,
+    register_local_autograd_function,
+)
 from spmd_types._mesh_axis import MeshAxis
 from spmd_types._reinterpret_mesh import (  # noqa: F401
     _format_arg_for_context,
@@ -53,8 +57,9 @@ from spmd_types._state import _current_mode, _set_current_mode, current_mesh
 from spmd_types._traceback import _filter_and_reraise
 from spmd_types._type_attr import get_local_type
 from spmd_types.runtime import (  # noqa: F401
-    _LOCAL_AUTOGRAD_FUNCTIONS,
+    _get_autograd_spmd_typecheck,
     _PARTITION_SPEC_ATTR,
+    _run_autograd_spmd_typecheck,
     _set_local_type,
     _set_partition_spec,
     _TRACE,
@@ -68,9 +73,9 @@ 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,
     trace,
 )
 from spmd_types.types import (
@@ -2454,10 +2459,9 @@ from spmd_types._state import is_type_checking  # noqa: E402, F401
 # Set of autograd.Function subclasses whose apply is known to be local-only
 # (i.e., each output element depends only on the corresponding input elements
 # so the standard element-wise type propagation rule is safe).
-# Set of autograd.Function subclasses registered with a typecheck_forward
-# staticmethod.  When .apply() is intercepted by __torch_function__,
-# typecheck_forward is called INSTEAD of .apply().  It should assert_type()
-# on inputs, call .apply() to execute, and assert_type() on the output.
+# Set of autograd.Function subclasses registered with custom typechecking.
+# Legacy typecheck_forward wrappers execute the function themselves, while
+# spmd_typecheck runs after the real apply call managed by the checker.
 
 
 # The original C++ descriptor for autograd.Function.apply, saved when the
@@ -2531,6 +2535,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 +2549,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 +2563,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 +2631,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,20 +2704,28 @@ 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
+            # Typecheck-registered autograd function. The mode is popped off the
             # TorchFunctionMode stack during __torch_function__ dispatch (via
-            # _pop_mode_temporarily in torch/overrides.py), so .apply() called
-            # inside typecheck_forward executes normally without re-entering
-            # this method.
+            # _pop_mode_temporarily in torch/overrides.py), so the real apply
+            # executes normally without re-entering this method.
             autograd_cls = _get_autograd_function_class(func)
-            if (
-                autograd_cls is not None
-                and autograd_cls in _TYPECHECK_AUTOGRAD_FUNCTIONS
-            ):
-                return autograd_cls.typecheck_forward(*args, **kwargs)
+            if autograd_cls is not None:
+                spmd_typecheck = _get_autograd_spmd_typecheck(autograd_cls)
+                if spmd_typecheck is not None:
+                    if kwargs:
+                        raise TypeError(
+                            f"{autograd_cls.__name__}.apply() accepts positional "
+                            "arguments only"
+                        )
+                    return _run_autograd_spmd_typecheck(autograd_cls, func, args)
+                if autograd_cls in _TYPECHECK_AUTOGRAD_FUNCTIONS:
+                    return autograd_cls.typecheck_forward(*args, **kwargs)
 
             # DTensor tracks its own placement metadata; the SPMD type checker
             # should not interfere with DTensor operations.  This covers
@@ -2744,6 +2777,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 +2830,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.
@@ -2859,10 +2915,9 @@ class _SpmdTypeMode(torch.overrides.TorchFunctionMode):
                         f"checking. Use "
                         f"{register_local_autograd_function.__name__}("
                         f"{autograd_cls.__name__}) to mark it as safe "
-                        f"for type propagation, and provide "
-                        f"register_autograd_function("
-                        f"{autograd_cls.__name__}) for a custom "
-                        f"typecheck_forward."
+                        f"for type propagation, or define "
+                        f"{autograd_cls.__name__}.spmd_typecheck for a custom "
+                        f"typecheck rule."
                     )
                 return result
 
@@ -2916,15 +2971,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 +3142,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.4/spmd_types/_local_registration.py

@@ -0,0 +1,50 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Lightweight registries for local SPMD type propagation."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+_LOCAL_AUTOGRAD_FUNCTIONS: set[type] = set()
+_LOCAL_BACKWARD_HOOKS: set[Callable] = set()
+
+
+def register_local_autograd_function(cls: type) -> type:
+    """Register an autograd.Function subclass as local-only for SPMD type checking.
+
+    Local-only means the function's forward operates element-wise (or more
+    generally, does not rearrange data across the tensor in a way that would
+    change its sharding type). It must not perform collectives or cross-device
+    communication. For functions that do, use ``register_autograd_function``
+    with custom typechecking instead.
+
+    Registered functions get the standard local type propagation rule when
+    type checking is active:
+
+    - Inputs may freely mix R and V types; the output is R unless any input
+      is V, in which case it is V.
+    - All-I inputs produce I outputs.
+    - R/V and I cannot be mixed.
+    - P is forbidden.
+
+    Unregistered autograd functions that reach the type checker leave their
+    outputs untyped or raise in strict mode.
+    """
+    _LOCAL_AUTOGRAD_FUNCTIONS.add(cls)
+    return cls
+
+
+def register_local_backward_hook(fn: Callable) -> Callable:
+    """Declare that ``fn`` does not alter SPMD types when it runs in backward.
+
+    Analogous to ``register_local_autograd_function``, the hook's backward is
+    treated as local with no collectives or type-changing effects on gradients.
+    This covers both full backward post-hooks and pre-hooks.
+    """
+    _LOCAL_BACKWARD_HOOKS.add(fn)
+    return fn

{spmd_types-0.2.2 → spmd_types-0.2.4}/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.4}/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.4}/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.4}/spmd_types/runtime.py

@@ -25,16 +25,20 @@ need the annotation APIs.
 from __future__ import annotations
 
 import builtins
+import inspect
 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._local_registration import (  # noqa: F401
+    register_local_autograd_function,
+)
 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 +593,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 +646,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)
 
 
@@ -659,43 +735,37 @@ def mutate_type(
 # Autograd function registration
 # =============================================================================
 
-# Sets of autograd.Function subclasses registered for type checking.
-# _LOCAL_AUTOGRAD_FUNCTIONS: element-wise / local-only functions.
-# _TYPECHECK_AUTOGRAD_FUNCTIONS: functions with custom typecheck_forward.
-_LOCAL_AUTOGRAD_FUNCTIONS: set[type] = set()
+# Autograd.Function subclasses with custom typechecking.
 _TYPECHECK_AUTOGRAD_FUNCTIONS: set[type] = set()
 
 
-def register_local_autograd_function(cls: type) -> type:
-    """Register an autograd.Function subclass as local-only for SPMD type checking.
-
-    Local-only means the function's forward operates element-wise (or more
-    generally, does not rearrange data across the tensor in a way that would
-    change its sharding type).  It must NOT perform any collectives or
-    cross-device communication.  For functions that do, use
-    :func:`register_autograd_function` with a ``typecheck_forward`` method instead.
-
-    Registered functions get the standard local type propagation rule when
-    type checking is active:
+def _run_autograd_spmd_typecheck(
+    cls: type,
+    func: Callable[..., object],
+    args: tuple[object, ...],
+) -> object:
+    signature = inspect.signature(cls.forward)
+    has_ctx = cls.setup_context is torch.autograd.Function.setup_context
+    bound = signature.bind(*(args if not has_ctx else (None, *args)))
+    bound.apply_defaults()
+    forward_args = dict(bound.arguments)
+    if has_ctx:
+        forward_args.pop(next(iter(signature.parameters)))
 
-    - Inputs may freely mix R and V types; the output is R unless any input
-      is V, in which case it is V.
-    - All-I inputs produce I outputs.
-    - R/V and I cannot be mixed.
-    - P is forbidden.
+    outputs = func(*args)
+    hook = cls.spmd_typecheck
+    names = tuple(inspect.signature(hook).parameters)[1:]
+    hook(outputs, **{name: forward_args[name] for name in names})
 
-    Unregistered autograd functions that reach the type checker will leave
-    their outputs untyped (or raise in strict mode), since the checker
-    cannot know whether the function is safe for automatic type propagation.
+    return outputs
 
-    Can be used as a decorator::
 
-        @register_local_autograd_function
-        class MyOp(torch.autograd.Function):
-            ...
-    """
-    _LOCAL_AUTOGRAD_FUNCTIONS.add(cls)
-    return cls
+def _get_autograd_spmd_typecheck(cls: type) -> Callable[..., object] | None:
+    """Return and validate an autograd.Function's SPMD typecheck hook."""
+    hook = getattr(cls, "spmd_typecheck", None)
+    if hook is not None and not callable(hook):
+        raise TypeError(f"{cls.__name__}.spmd_typecheck must be callable")
+    return hook
 
 
 def register_autograd_function(cls: type) -> type:
@@ -705,34 +775,44 @@ def register_autograd_function(cls: type) -> type:
     non-trivial type transformations where the default local-only rule
     would produce incorrect output types.
 
-    The class must define a ``typecheck_forward`` staticmethod that receives
-    the same positional/keyword arguments as ``.apply()``.  Inside, it should
-    call ``assert_type`` on inputs, call ``.apply()`` to run the function,
-    then call ``assert_type`` on the output.  This is symmetric with
-    ``typecheck_forward`` on ``nn.Module`` subclasses in llama4x::
+    A class may define an ``spmd_typecheck`` staticmethod. Its presence is
+    detected automatically, so new classes do not need this decorator. The
+    hook runs after the function, receives its exact return value as the first
+    argument, and names only the forward arguments it needs::
 
-        @register_autograd_function
         class MyCollectiveOp(torch.autograd.Function):
             @staticmethod
             def forward(ctx, x, y):
                 return x + y
 
             @staticmethod
-            def typecheck_forward(x, y):
+            def spmd_typecheck(outputs, *, x):
                 assert_type(x, {pg: S(-1)})
-                out = MyCollectiveOp.apply(x, y)
-                assert_type(out, {pg: R})
-                return out
+                assert_type_like(outputs, x, {pg: R})
 
             @staticmethod
             def backward(ctx, g):
                 return g, g
+
+    The legacy ``typecheck_forward`` wrapper remains supported for existing
+    registrations, but cannot be combined with ``spmd_typecheck``.
     """
-    if not callable(getattr(cls, "typecheck_forward", None)):
+    legacy_hook = getattr(cls, "typecheck_forward", None)
+    spmd_typecheck = _get_autograd_spmd_typecheck(cls)
+
+    if legacy_hook is not None and spmd_typecheck is not None:
+        raise TypeError(
+            f"{cls.__name__} cannot define typecheck_forward together with "
+            "spmd_typecheck"
+        )
+    if legacy_hook is None and spmd_typecheck is None:
         raise TypeError(
-            f"{cls.__name__} must define a typecheck_forward staticmethod "
+            f"{cls.__name__} must define typecheck_forward or spmd_typecheck "
             f"when using @register_autograd_function"
         )
+    if legacy_hook is not None and not callable(legacy_hook):
+        raise TypeError(f"{cls.__name__}.typecheck_forward must be callable")
+
     _TYPECHECK_AUTOGRAD_FUNCTIONS.add(cls)
     return cls
 
@@ -908,7 +988,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 +1085,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 +1203,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.4/spmd_types.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: spmd_types
-Version: 0.2.2
+Version: 0.2.4
 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.4}/spmd_types.egg-info/SOURCES.txt

@@ -10,6 +10,7 @@ spmd_types/_dtensor_checker.py
 spmd_types/_dtype_utils.py
 spmd_types/_frame.py
 spmd_types/_local.py
+spmd_types/_local_registration.py
 spmd_types/_mesh.py
 spmd_types/_mesh_axis.py
 spmd_types/_mesh_region.py
@@ -29,6 +30,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