spmd_types 0.2.3__tar.gz → 0.2.4__tar.gz

{spmd_types-0.2.3/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.3
+Version: 0.2.4
 Summary: A type system for distributed (SPMD) tensor computations in PyTorch
 Author: Meta Platforms, Inc.
 License: BSD 3-Clause License

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

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "spmd_types"
-version = "0.2.3"
+version = "0.2.4"
 description = "A type system for distributed (SPMD) tensor computations in PyTorch"
 readme = "README.md"
 license = {file = "LICENSE"}

{spmd_types-0.2.3 → 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
 
@@ -64,7 +65,6 @@ from spmd_types.runtime import (  # noqa: F401
     no_typecheck,
     register_autograd_function,
     register_decomposition,
-    register_local_autograd_function,
     trace,
 )
 

{spmd_types-0.2.3 → 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.3 → 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,
@@ -71,7 +76,6 @@ from spmd_types.runtime import (  # noqa: F401
     no_typecheck,
     register_autograd_function,
     register_decomposition,
-    register_local_autograd_function,
     trace,
 )
 from spmd_types.types import (
@@ -2455,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
@@ -2707,18 +2710,22 @@ class _SpmdTypeMode(torch.overrides.TorchFunctionMode):
                         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
@@ -2908,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
 

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

@@ -25,6 +25,7 @@ need the annotation APIs.
 from __future__ import annotations
 
 import builtins
+import inspect
 import logging
 import os
 from contextlib import AbstractContextManager, contextmanager
@@ -32,6 +33,9 @@ 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 _no_typecheck_context, current_mesh, is_type_checking
@@ -731,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:
@@ -777,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__} must define a typecheck_forward staticmethod "
+            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 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
 

{spmd_types-0.2.3 → 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.3
+Version: 0.2.4
 Summary: A type system for distributed (SPMD) tensor computations in PyTorch
 Author: Meta Platforms, Inc.
 License: BSD 3-Clause License

{spmd_types-0.2.3 → 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