spmd_types 0.1.0__tar.gz

spmd_types-0.1.0/LICENSE

@@ -0,0 +1,28 @@
+BSD 3-Clause License
+
+Copyright (c) 2025, Edward Z. Yang
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

spmd_types-0.1.0/PKG-INFO

@@ -0,0 +1,19 @@
+Metadata-Version: 2.4
+Name: spmd_types
+Version: 0.1.0
+Summary: A type system for distributed (SPMD) tensor computations in PyTorch
+Author-email: "Edward Z. Yang" <ezyang@mit.edu>
+License-Expression: BSD-3-Clause
+Project-URL: Homepage, https://github.com/ezyang/spmd_types
+Project-URL: Repository, https://github.com/ezyang/spmd_types
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+License-File: LICENSE
+Requires-Dist: torch
+Dynamic: license-file

spmd_types-0.1.0/pyproject.toml

@@ -0,0 +1,36 @@
+[build-system]
+requires = ["setuptools>=68.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "spmd_types"
+version = "0.1.0"
+description = "A type system for distributed (SPMD) tensor computations in PyTorch"
+requires-python = ">=3.10"
+license = "BSD-3-Clause"
+authors = [
+    { name = "Edward Z. Yang", email = "ezyang@mit.edu" },
+]
+dependencies = [
+    "torch",
+]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+
+[project.urls]
+Homepage = "https://github.com/ezyang/spmd_types"
+Repository = "https://github.com/ezyang/spmd_types"
+
+[tool.setuptools.packages.find]
+include = ["sixlib*"]
+exclude = ["sixlib.spmd_types.*_test*", "sixlib.spmd_types.eval*"]
+
+[tool.setuptools.package-data]
+"sixlib.spmd_types" = ["py.typed"]

spmd_types-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

spmd_types-0.1.0/sixlib/spmd_types/__init__.py

@@ -0,0 +1,92 @@
+# spmd_types package
+from __future__ import annotations
+
+import sys
+from types import ModuleType
+
+# Collectives and operations
+from sixlib.spmd_types._checker import (  # noqa: F401
+    _SpmdTypeBackwardCompatibleMode as SpmdTypeMode,
+    assert_local_type,
+    assert_type,
+    is_type_checking,
+    mutate_type,
+    no_typecheck,
+    register_autograd_function,
+    register_local_autograd_function,
+    trace,
+    typecheck,
+)
+from sixlib.spmd_types._collectives import (  # noqa: F401
+    all_gather,
+    all_reduce,
+    all_to_all,
+    redistribute,
+    reduce_scatter,
+    unshard,
+)
+from sixlib.spmd_types._dist import set_dist  # noqa: F401
+from sixlib.spmd_types._dtensor import (  # noqa: F401
+    dtensor_placement_to_spmd_type,
+    spmd_redistribute,
+    spmd_type_to_dtensor_placement,
+)
+from sixlib.spmd_types._local import (  # noqa: F401
+    convert,
+    invariant_to_replicate,
+    reinterpret,
+    shard,
+)
+from sixlib.spmd_types._mesh_axis import MeshAxis  # noqa: F401
+from sixlib.spmd_types._scalar import Scalar  # noqa: F401
+from sixlib.spmd_types._traceback import traceback_filtering  # noqa: F401
+from sixlib.spmd_types._type_attr import get_local_type  # noqa: F401
+
+# Types
+from sixlib.spmd_types.types import (  # noqa: F401
+    DimSharding,
+    I,
+    Invariant,
+    LocalSpmdType,
+    normalize_axis,
+    P,
+    Partial,
+    PartitionSpec,
+    PerMeshAxisLocalSpmdType,
+    PerMeshAxisSpmdType,
+    PerMeshAxisSpmdTypes,
+    R,
+    Replicate,
+    S,
+    Shard,
+    SpmdTypeError,
+    TensorSharding,
+    V,
+    Varying,
+)
+
+
+class _SpmdTypesModule(ModuleType):
+    """Module wrapper that provides TYPE_CHECKING as a dynamic attribute.
+
+    Accessing ``sixlib.spmd_types.TYPE_CHECKING`` returns True when a
+    ``typecheck()`` context is active and False otherwise.  This uses the
+    sys.modules replacement trick (same pattern as torch._dynamo.config)
+    so that a simple attribute read works like a function call.
+    """
+
+    @property
+    def TYPE_CHECKING(self) -> bool:
+        return is_type_checking()
+
+
+# Replace this module in sys.modules so that attribute access on the
+# module object goes through _SpmdTypesModule.__getattr__ / properties.
+_self = sys.modules[__name__]
+_obj = _SpmdTypesModule(__name__)
+_obj.__dict__.update(_self.__dict__)
+_obj.__file__ = _self.__file__
+_obj.__package__ = _self.__package__
+_obj.__path__ = _self.__path__  # type: ignore[attr-defined]
+_obj.__spec__ = _self.__spec__
+sys.modules[__name__] = _obj