torchdt 0.1.0__tar.gz

torchdt-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 XLNS
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

torchdt-0.1.0/PKG-INFO

@@ -0,0 +1,35 @@
+Metadata-Version: 2.4
+Name: torchdt
+Version: 0.1.0
+Summary: torchdt is a python package that provides custom datatypes for training and inference in PyTorch.
+Author-email: Alexander Krentz <alexkrentz2@gmail.com>
+License: MIT License
+        
+        Copyright (c) 2024 XLNS
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+Project-URL: Homepage, https://github.com/akrentz6/torchdt
+Keywords: python,pytorch,machine learning
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.0
+Requires-Dist: numpy
+Dynamic: license-file
+
+A Python addon package to PyTorch that provides experimentation with custom tensor datatypes for training and inference.

torchdt-0.1.0/README.md

@@ -0,0 +1 @@
+A Python addon package to PyTorch that provides experimentation with custom tensor datatypes for training and inference.

torchdt-0.1.0/pyproject.toml

@@ -0,0 +1,41 @@
+[build-system]
+requires = [
+    "setuptools>=69.0",
+    "wheel",
+    "torch>=2.0",
+    "ninja",
+]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "torchdt"
+version = "0.1.0"
+description = "torchdt is a python package that provides custom datatypes for training and inference in PyTorch."
+readme = "README.md"
+license = {file = "LICENSE"}
+authors = [{name = "Alexander Krentz", email = "alexkrentz2@gmail.com"}]
+keywords = ["python", "pytorch", "machine learning"]
+dependencies = ["torch>=2.0", "numpy"]
+
+[project.urls]
+"Homepage" = "https://github.com/akrentz6/torchdt"
+
+[tool.setuptools]
+include-package-data = true
+
+[tool.setuptools.package-data]
+"torchdt" = [
+    "include/torchdt/**/*.h",
+    "csrc/**/*.cpp",
+    "csrc/**/*.h",
+]
+
+[[tool.setuptools.ext-modules]]
+name = "torchdt._C"
+sources = []  # Filled dynamically
+language = "c++"
+extra-compile-args = ["-O3", "-std=c++17"]
+optional = true
+
+[tool.setuptools.cmdclass]
+build_ext = "torchdt._build_ext.BuildCxxExtension"

torchdt-0.1.0/setup.cfg

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

torchdt-0.1.0/torchdt/__init__.py

@@ -0,0 +1,11 @@
+from . import transforms
+from .datatype import DType
+from . import lns
+from . import ops
+from . import autograd
+from . import funcs
+from . import optim
+
+__all__ = [
+    "DType",
+]

torchdt-0.1.0/torchdt/_build_ext.py

@@ -0,0 +1,71 @@
+from torch.utils.cpp_extension import BuildExtension, include_paths
+from pathlib import Path
+import warnings
+import os
+
+# The exact exception raised by a compiler/linker failure can vary
+# between setuptools versions, so we catch the whole family.
+try:
+    # setuptools ≥ 68
+    from setuptools.errors import (
+        CompileError,
+        LinkError,
+        PlatformError,
+    )
+except ImportError:
+    # setuptools < 68
+    from distutils.errors import ( # type: ignore
+        CompileError,
+        LinkError,
+        DistutilsPlatformError as PlatformError,
+    )
+
+class BuildCxxExtension(BuildExtension):
+    """
+    Build all .cpp files in csrc/ and include Torch headers.
+    Fails hard if compilation fails.
+    """
+
+    user_options = BuildExtension.user_options + [
+        ("no-cpp", None, "Do not build the C++ extension"),
+    ]
+    boolean_options = getattr(BuildExtension, "boolean_options", []) + ["no-cpp"]
+
+    def initialize_options(self):
+        super().initialize_options()
+        self.no_cpp = False
+
+    def finalize_options(self):
+        super().finalize_options()
+        if os.environ.get("TORCHDT_NO_CPP", "").lower() in {"1", "true", "yes", "on"}:
+            self.no_cpp = True
+
+    def run(self):
+        try:
+            super().run()
+        except (CompileError, LinkError, PlatformError) as exc:
+            warnings.warn(
+                "C++ backend for xlnstorch was not built; "
+                "installation will succeed but will run in pure-python mode.",
+                RuntimeWarning,
+            )
+
+    def build_extensions(self):
+        if self.no_cpp:
+            self.extensions = []
+            return
+
+        src_dir = Path(__file__).resolve().parent / "csrc"
+        include_dir = Path(__file__).resolve().parent / "include"
+        cpp_files = [str(p) for p in src_dir.glob("*.cpp")]
+
+        torch_includes = include_paths()
+
+        for ext in self.extensions:
+            if ext.name == "torchdt._C":
+                ext.sources = cpp_files
+            ext.include_dirs = list(ext.include_dirs or [])
+            ext.include_dirs += torch_includes
+            ext.include_dirs += [str(include_dir)]
+
+        super().build_extensions()

torchdt-0.1.0/torchdt/autograd.py

@@ -0,0 +1,280 @@
+import torch
+from torch.autograd.graph import get_gradient_edge
+import functools
+from typing import Optional, Tuple
+
+__all__ = [
+    "DTFunction",
+    "DTNonDifferentiableFunction",
+]
+
+def _find_first_grad_tensor(args):
+    if not isinstance(args, (list, tuple)):
+        args = (args,)
+    for arg in args:
+        if isinstance(arg, torch.Tensor) and arg.requires_grad:
+            return arg
+    return None
+
+def _cast_int(x, dtype):
+    return x.view(dtype.int_dtype) if isinstance(x, torch.Tensor) and x.dtype == dtype.float_dtype else x
+
+def _cast_float(x, dtype):
+    return x.view(dtype.float_dtype) if isinstance(x, torch.Tensor) and x.dtype == dtype.int_dtype else x
+
+def _cast_values(values, indices, cast_fn, dtype):
+    all_indices = indices is None
+
+    if isinstance(values, tuple):
+        return tuple(
+            cast_fn(values[i], dtype) if all_indices or i in indices else values[i]
+            for i in range(len(values))
+        )
+
+    elif isinstance(values, list):
+        return [
+            cast_fn(values[i], dtype) if all_indices or i in indices else values[i]
+            for i in range(len(values))
+        ]
+
+    else:
+        return cast_fn(values, dtype) if all_indices or 0 in indices else values
+
+# added to forward only if setup_context is not defined
+def forward_ctx_decorator(func, cls):
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        ctx = args[0]
+        ops = args[1]
+        inputs = args[2:]
+
+        ctx._dtype = cls.dtype
+        ctx._input_indices = cls.input_indices
+
+        cast_inputs = _cast_values(inputs, cls.input_indices, _cast_int, cls.dtype)
+        output = func(ctx, ops, *cast_inputs, **kwargs)
+        return _cast_values(output, cls.output_indices, _cast_float, cls.dtype)
+
+    return wrapper
+
+def forward_no_ctx_decorator(func, cls):
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        ops = args[0]
+        inputs = args[1:]
+
+        cast_inputs = _cast_values(inputs, cls.input_indices, _cast_int, cls.dtype)
+        output = func(ops, *cast_inputs, **kwargs)
+        return _cast_values(output, cls.output_indices, _cast_float, cls.dtype)
+
+    return wrapper
+
+def setup_context_decorator(func, cls):
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        ctx = args[0]
+        ops = args[1][0]
+        inputs = args[1][1:]
+        output = args[2]
+
+        ctx._dtype = cls.dtype
+        ctx._input_indices = cls.input_indices
+
+        cast_inputs = _cast_values(inputs, cls.input_indices, _cast_int, cls.dtype)
+        cast_output = _cast_values(output, cls.output_indices, _cast_int, cls.dtype)
+        return func(ctx, ops, cast_inputs, cast_output, **kwargs)
+
+    return wrapper
+
+def backward_decorator(func, cls):
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        ctx = args[0]
+        grads = args[1:]
+
+        dtype = ctx._dtype
+        ops = ctx._dtype.ops
+        input_indices = ctx._input_indices
+
+        cast_grads = _cast_values(grads, cls.output_indices, _cast_int, dtype)
+        output = func(ctx, ops, *cast_grads, **kwargs)
+        cast_output = _cast_values(output, input_indices, _cast_float, dtype)
+
+        if isinstance(output, tuple):
+            return (None,) + cast_output
+        return None, cast_output
+
+    return wrapper
+
+class DTFunction(torch.autograd.Function):
+    """
+    Parent class for custom autograd Functions that work with DType tensors.
+    Subclasses should implement static methods `forward` and `backward` (and
+    optionally `setup_context`).
+    """
+
+    output_indices: Optional[Tuple[int, ...]] = None
+
+    @classmethod
+    def _register_ops_decorator(cls, func_name, decorator):
+        func = getattr(cls, func_name)
+
+        if not getattr(func, "_is_decorated", False):
+            decorated = decorator(func, cls)
+            decorated._is_lns_decorated = True
+            setattr(cls, func_name, decorated)
+
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+
+        # check if setup_context is defined in subclass or inherited
+        # from torch.autograd.function._SingleLevelFunction. If it is
+        # inherited, we only need to decorate forward and backward.
+        # If it is defined, we need to decorate setup_context instead
+        # of forward.
+        if cls.setup_context is torch.autograd.function._SingleLevelFunction.setup_context:
+            cls._register_ops_decorator("forward", forward_ctx_decorator)
+        else:
+            cls._register_ops_decorator("forward", forward_no_ctx_decorator)
+            cls._register_ops_decorator("setup_context", setup_context_decorator)
+
+        cls._register_ops_decorator("backward", backward_decorator)
+
+    @classmethod
+    def apply(cls, *args, **kwargs):
+        from torchdt import DType # avoid circular import
+
+        if kwargs:
+            raise ValueError(
+                "torch.autograd.Function does not support keyword arguments. "
+                "Please use positional arguments only."
+            )
+
+        input_indices = []
+        subtypes = []
+        prepped_inputs = []
+
+        for i, arg in enumerate(args):
+            if isinstance(arg, DType):
+                subtypes.append(arg.__class__)
+                input_indices.append(i)
+                prepped_inputs.append(arg._float)
+            else:
+                prepped_inputs.append(arg)
+
+        if not subtypes:
+            raise ValueError("DTFunction.apply() requires at least one DType tensor argument.")
+
+        dtype = subtypes[0]
+        if any(st != dtype for st in subtypes):
+            raise ValueError("All DType arguments to DTFunction.apply() must be of the same type.")
+
+        # perform the operation using the Ops class for this DType
+        cls.dtype = dtype
+        cls.input_indices = input_indices
+        result = super().apply(dtype.ops, *prepped_inputs)
+        del cls.input_indices
+        del cls.dtype
+
+        # get gradient edge to correctly handle grads for DType tensors
+        first_tensor = _find_first_grad_tensor(result)
+        if first_tensor is not None:
+            edge = get_gradient_edge(first_tensor)
+
+            j = 0
+            for arg in args:
+
+                # ignore non-tensor arguments
+                if not isinstance(arg, torch.Tensor):
+                    continue
+                j += 1
+
+                # only register hooks for DType inputs with gradients
+                if not (isinstance(arg, dtype) and arg.requires_grad):
+                    continue
+
+                arg._track_operation(edge, j - 1)
+
+        if isinstance(result, torch.Tensor):
+            return dtype(
+                result, internal=True
+            ) if cls.output_indices is None or 0 in cls.output_indices else result
+
+        elif isinstance(result, list):
+            return [
+                dtype(result[i], internal=True)
+                if cls.output_indices is None or i in cls.output_indices else result[i]
+                for i in range(len(result))
+            ]
+
+        elif isinstance(result, tuple):
+            return tuple(
+                dtype(result[i], internal=True)
+                if cls.output_indices is None or i in cls.output_indices else result[i]
+                for i in range(len(result))
+            )
+
+        return result
+
+class DTNonDifferentiableFunction:
+
+    output_indices: Optional[Tuple[int, ...]] = None
+
+    @staticmethod
+    def forward(ops, *args, **kwargs):
+        """Forward pass for non-differentiable function."""
+        raise NotImplementedError
+
+    @classmethod
+    def apply(cls, *args, **kwargs):
+        from torchdt import DType # avoid circular import
+
+        if kwargs:
+            raise ValueError(
+                "DTNonDifferentiableFunction does not support keyword arguments. "
+                "Please use positional arguments only."
+            )
+
+        subtypes = []
+        prepped_inputs = []
+
+        for i, arg in enumerate(args):
+            if isinstance(arg, DType):
+                subtypes.append(arg.__class__)
+                prepped_inputs.append(arg._int)
+            else:
+                prepped_inputs.append(arg)
+
+        if not subtypes:
+            raise ValueError("DTFunction.apply() requires at least one DType tensor argument.")
+
+        dtype = subtypes[0]
+        if any(st != dtype for st in subtypes):
+            raise ValueError("All DType arguments to DTFunction.apply() must be of the same type.")
+
+        result = cls.forward(dtype.ops, *prepped_inputs)
+
+        if isinstance(result, torch.Tensor):
+            return dtype(
+                result, internal=True
+            ) if cls.output_indices is None or 0 in cls.output_indices else result
+
+        elif isinstance(result, list):
+            return [
+                dtype(result[i], internal=True)
+                if cls.output_indices is None or i in cls.output_indices else result[i]
+                for i in range(len(result))
+            ]
+
+        elif isinstance(result, tuple):
+            return tuple(
+                dtype(result[i], internal=True)
+                if cls.output_indices is None or i in cls.output_indices else result[i]
+                for i in range(len(result))
+            )
+
+        return result

torchdt-0.1.0/torchdt/csrc/dispatcher.h

@@ -0,0 +1,101 @@
+#ifndef DISPATCHER_H
+#define DISPATCHER_H
+
+#include <torchdt/registry.h>
+#include "ops_impl.h"
+
+class Dispatcher {
+public:
+    static Dispatcher& instance() {
+        static Dispatcher d;
+        return d;
+    }
+
+    // Get cached OpsBase* for name and bitwidth.
+    // If not cached and the public registry contains an Ops<bitwidth> for this name,
+    // lazily construct an OpsImpl<bitwidth> from the stored Ops and cache it.
+    // Returns nullptr if no Ops registered for (name,bitwidth).
+    OpsBase* get(const std::string& name, size_t bitwidth) {
+        const std::string key = make_key(name, bitwidth);
+
+        // fast path: check cached map under lock
+        {
+            std::lock_guard<std::mutex> guard(mutex_);
+            auto it = impl_map_.find(key);
+            if (it != impl_map_.end()) return it->second.get();
+        }
+
+        // not cached: try to construct based on bitwidth
+        std::unique_ptr<OpsBase> created = make_impl_from_registry(name, bitwidth);
+        if (!created) return nullptr;
+
+        // store & return
+        {
+            std::lock_guard<std::mutex> guard(mutex_);
+            // double-check another thread didn't create it
+            auto it = impl_map_.find(key);
+            if (it == impl_map_.end()) {
+                OpsBase* raw = created.get();
+                impl_map_.emplace(key, std::move(created));
+                return raw;
+            } else {
+                return it->second.get();
+            }
+        }
+    }
+
+    // Convenience template: typed get that returns pointer to OpsImpl<bitwidth>
+    template <size_t bitwidth>
+    OpsImpl<bitwidth>* get_typed(const std::string& name) {
+        OpsBase* base = get(name, bitwidth);
+        return static_cast<OpsImpl<bitwidth>*>(base);
+    }
+
+private:
+    Dispatcher() = default;
+
+    // Create OpsImpl for supported bitwidths by fetching the registered Ops<bitwidth> from public registry
+    std::unique_ptr<OpsBase> make_impl_from_registry(const std::string& name, size_t bitwidth) {
+        switch (bitwidth) {
+        case 8: {
+            auto p = Registry::instance().get_ops_typed<8>(name);
+            if (!p) return nullptr;
+            return std::make_unique<OpsImpl<8>>(*p);
+        }
+        case 16: {
+            auto p = Registry::instance().get_ops_typed<16>(name);
+            if (!p) return nullptr;
+            return std::make_unique<OpsImpl<16>>(*p);
+        }
+        case 32: {
+            auto p = Registry::instance().get_ops_typed<32>(name);
+            if (!p) return nullptr;
+            return std::make_unique<OpsImpl<32>>(*p);
+        }
+        case 64: {
+            auto p = Registry::instance().get_ops_typed<64>(name);
+            if (!p) return nullptr;
+            return std::make_unique<OpsImpl<64>>(*p);
+        }
+        default:
+            return nullptr;
+        }
+    }
+
+    std::map<std::string, std::unique_ptr<OpsBase>> impl_map_;
+    std::mutex mutex_;
+};
+
+// -----------------------------
+// Helper functions (optional)
+// -----------------------------
+inline OpsBase* get_ops_impl(const std::string& name, size_t bitwidth) {
+    return Dispatcher::instance().get(name, bitwidth);
+}
+
+template <size_t bitwidth>
+inline OpsImpl<bitwidth>* get_ops_impl_typed(const std::string& name) {
+    return Dispatcher::instance().get_typed<bitwidth>(name);
+}
+
+#endif // DISPATCHER_H