nshutils 0.1.0__tar.gz

nshutils-0.1.0/PKG-INFO

@@ -0,0 +1,21 @@
+Metadata-Version: 2.1
+Name: nshutils
+Version: 0.1.0
+Summary: 
+Author: Nima Shoghi
+Author-email: nimashoghi@gmail.com
+Requires-Python: >=3.10,<4.0
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: beartype (>=0.18.5,<0.19.0)
+Requires-Dist: jaxtyping (>=0.2.33,<0.3.0)
+Requires-Dist: lovely-jax (>=0.1.3,<0.2.0)
+Requires-Dist: lovely-numpy (>=0.2.13,<0.3.0)
+Requires-Dist: lovely-tensors (>=0.1.16,<0.2.0)
+Requires-Dist: pysnooper (>=1.2.0,<2.0.0)
+Requires-Dist: typing-extensions
+Description-Content-Type: text/markdown
+
+

nshutils-0.1.0/pyproject.toml

@@ -0,0 +1,21 @@
+[tool.poetry]
+name = "nshutils"
+version = "0.1.0"
+description = ""
+authors = ["Nima Shoghi <nimashoghi@gmail.com>"]
+readme = "README.md"
+
+[tool.poetry.dependencies]
+python = "^3.10"
+pysnooper = "^1.2.0"
+jaxtyping = "^0.2.33"
+typing-extensions = "*"
+lovely-numpy = {version = "^0.2.13", optional = true}
+lovely-tensors = {version = "^0.1.16", optional = true}
+beartype = "^0.18.5"
+lovely-jax = {version = "^0.1.3", optional = true}
+
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"

nshutils-0.1.0/src/nshutils/__init__.py

@@ -0,0 +1,2 @@
+from . import typecheck as typecheck
+from .snoop import snoop as snoop

nshutils-0.1.0/src/nshutils/snoop.py

@@ -0,0 +1,216 @@
+import contextlib
+from typing import Any, Protocol, cast
+
+from typing_extensions import TypeVar
+
+T = TypeVar("T", infer_variance=True)
+
+
+class SnoopConstructor(Protocol):
+    def __call__(self, *args, **kwargs) -> contextlib.AbstractContextManager: ...
+
+    def disable(self) -> contextlib.AbstractContextManager: ...
+
+
+try:
+    import warnings
+    from contextlib import nullcontext
+
+    import lovely_numpy as lo
+    import lovely_tensors as lt
+    import numpy
+    import pysnooper
+    import pysnooper.utils
+    import torch
+    from pkg_resources import DistributionNotFound, get_distribution
+
+    FLOATING_POINTS = set()
+    for i in ["float", "double", "half", "complex128", "complex32", "complex64"]:
+        if hasattr(torch, i):  # older version of PyTorch do not have complex dtypes
+            FLOATING_POINTS.add(getattr(torch, i))
+
+    try:
+        __version__ = get_distribution(__name__).version
+    except DistributionNotFound:
+        # package is not installed
+        pass
+
+    def default_format(x):
+        try:
+            formatted = str(lt.lovely(x))
+            return formatted
+        except BaseException:
+            return str(x.shape)
+
+    def default_numpy_format(x):
+        return str(lo.lovely(x))
+
+    class TorchSnooper(pysnooper.tracer.Tracer):
+        def __init__(
+            self,
+            *args,
+            tensor_format=default_format,
+            numpy_format=default_numpy_format,
+            **kwargs,
+        ):
+            self.orig_custom_repr = (
+                kwargs["custom_repr"] if "custom_repr" in kwargs else ()
+            )
+            custom_repr = (lambda x: True, self.compute_repr)
+            kwargs["custom_repr"] = (custom_repr,)
+            super(TorchSnooper, self).__init__(*args, **kwargs)
+            self.tensor_format = tensor_format
+            self.numpy_format = numpy_format
+
+        @staticmethod
+        def is_return_types(x):
+            return type(x).__module__ == "torch.return_types"
+
+        def return_types_repr(self, x):
+            if type(x).__name__ in {
+                "max",
+                "min",
+                "median",
+                "mode",
+                "sort",
+                "topk",
+                "kthvalue",
+            }:
+                return (
+                    type(x).__name__
+                    + "(values="
+                    + self.tensor_format(x.values)
+                    + ", indices="
+                    + self.tensor_format(x.indices)
+                    + ")"
+                )
+            if type(x).__name__ == "svd":
+                return (
+                    "svd(U="
+                    + self.tensor_format(x.U)
+                    + ", S="
+                    + self.tensor_format(x.S)
+                    + ", V="
+                    + self.tensor_format(x.V)
+                    + ")"
+                )
+            if type(x).__name__ == "slogdet":
+                return (
+                    "slogdet(sign="
+                    + self.tensor_format(x.sign)
+                    + ", logabsdet="
+                    + self.tensor_format(x.logabsdet)
+                    + ")"
+                )
+            if type(x).__name__ == "qr":
+                return (
+                    "qr(Q="
+                    + self.tensor_format(x.Q)
+                    + ", R="
+                    + self.tensor_format(x.R)
+                    + ")"
+                )
+            if type(x).__name__ == "solve":
+                return (
+                    "solve(solution="
+                    + self.tensor_format(x.solution)
+                    + ", LU="
+                    + self.tensor_format(x.LU)
+                    + ")"
+                )
+            if type(x).__name__ == "geqrf":
+                return (
+                    "geqrf(a="
+                    + self.tensor_format(x.a)
+                    + ", tau="
+                    + self.tensor_format(x.tau)
+                    + ")"
+                )
+            if type(x).__name__ in {"symeig", "eig"}:
+                return (
+                    type(x).__name__
+                    + "(eigenvalues="
+                    + self.tensor_format(x.eigenvalues)
+                    + ", eigenvectors="
+                    + self.tensor_format(x.eigenvectors)
+                    + ")"
+                )
+            if type(x).__name__ == "triangular_solve":
+                return (
+                    "triangular_solve(solution="
+                    + self.tensor_format(x.solution)
+                    + ", cloned_coefficient="
+                    + self.tensor_format(x.cloned_coefficient)
+                    + ")"
+                )
+            if type(x).__name__ == "gels":
+                return (
+                    "gels(solution="
+                    + self.tensor_format(x.solution)
+                    + ", QR="
+                    + self.tensor_format(x.QR)
+                    + ")"
+                )
+            warnings.warn("Unknown return_types encountered, open a bug report!")
+
+        def compute_repr(self, x):
+            orig_repr_func = pysnooper.utils.get_repr_function(x, self.orig_custom_repr)
+            if torch.is_tensor(x):
+                return self.tensor_format(x)
+            if isinstance(x, numpy.ndarray):
+                return self.numpy_format(x)
+            if self.is_return_types(x):
+                return self.return_types_repr(x)
+            if orig_repr_func is not repr:
+                return orig_repr_func(x)
+            if isinstance(x, (list, tuple)):
+                content = ""
+                for i in x:
+                    if content != "":
+                        content += ", "
+                    content += self.compute_repr(i)
+                if isinstance(x, tuple) and len(x) == 1:
+                    content += ","
+                if isinstance(x, tuple):
+                    return "(" + content + ")"
+                return "[" + content + "]"
+            if isinstance(x, dict):
+                content = ""
+                for k, v in x.items():
+                    if content != "":
+                        content += ", "
+                    content += self.compute_repr(k) + ": " + self.compute_repr(v)
+                return "{" + content + "}"
+            return repr(x)
+
+    class _Snoop:
+        disable = nullcontext
+        __call__ = TorchSnooper
+
+    snoop: SnoopConstructor = cast(Any, _Snoop())
+
+except ImportError:
+    import warnings
+    from contextlib import nullcontext
+
+    from typing_extensions import override
+
+    _has_warned = False
+
+    class _snoop_cls(nullcontext):
+        @classmethod
+        def disable(cls):
+            return nullcontext()
+
+        @override
+        def __enter__(self):
+            global _has_warned
+            if not _has_warned:
+                warnings.warn(
+                    "snoop is not installed, please install it to enable snoop"
+                )
+                _has_warned = True
+
+            return super().__enter__()
+
+    snoop: SnoopConstructor = cast(Any, _snoop_cls)

nshutils-0.1.0/src/nshutils/typecheck.py

@@ -0,0 +1,155 @@
+import os
+from collections.abc import Sequence
+from logging import getLogger
+from typing import Any
+
+from jaxtyping import BFloat16 as BFloat16
+from jaxtyping import Bool as Bool
+from jaxtyping import Complex as Complex
+from jaxtyping import Complex64 as Complex64
+from jaxtyping import Complex128 as Complex128
+from jaxtyping import Float as Float
+from jaxtyping import Float16 as Float16
+from jaxtyping import Float32 as Float32
+from jaxtyping import Float64 as Float64
+from jaxtyping import Inexact as Inexact
+from jaxtyping import Int as Int
+from jaxtyping import Int4 as Int4
+from jaxtyping import Int8 as Int8
+from jaxtyping import Int16 as Int16
+from jaxtyping import Int32 as Int32
+from jaxtyping import Int64 as Int64
+from jaxtyping import Integer as Integer
+from jaxtyping import Key as Key
+from jaxtyping import Num as Num
+from jaxtyping import Real as Real
+from jaxtyping import Shaped as Shaped
+from jaxtyping import UInt as UInt
+from jaxtyping import UInt4 as UInt4
+from jaxtyping import UInt8 as UInt8
+from jaxtyping import UInt16 as UInt16
+from jaxtyping import UInt32 as UInt32
+from jaxtyping import UInt64 as UInt64
+from jaxtyping._storage import get_shape_memo, shape_str
+from typing_extensions import TypeVar
+
+try:
+    import torch
+except ImportError:
+    torch = None
+
+try:
+    import np
+except ImportError:
+    np = None
+
+
+try:
+    import jax
+except ImportError:
+    jax = None
+log = getLogger(__name__)
+
+DISABLE_ENV_KEY = "NSHUTILS_DISABLE_TYPECHECKING"
+
+
+def typecheck_modules(modules: Sequence[str]):
+    """
+    Typecheck the given modules using `jaxtyping`.
+
+    Args:
+        modules: Modules to typecheck.
+    """
+    # If `DISABLE_ENV_KEY` is set and the environment variable is set, skip
+    #   typechecking.
+    if DISABLE_ENV_KEY is not None and bool(int(os.environ.get(DISABLE_ENV_KEY, "0"))):
+        log.critical(
+            f"Type checking is disabled due to the environment variable {DISABLE_ENV_KEY}."
+        )
+        return
+
+    # Install the jaxtyping import hook for this module.
+    from jaxtyping import install_import_hook
+
+    install_import_hook(modules, "beartype.beartype")
+
+    log.critical(f"Type checking the following modules: {modules}")
+
+
+def typecheck_this_module(additional_modules: Sequence[str] = ()):
+    """
+    Typecheck the calling module and any additional modules using `jaxtyping`.
+
+    Args:
+        additional_modules: Additional modules to typecheck.
+    """
+    # Get the calling module's name.
+    # Here, we can just use beartype's internal implementation behind
+    # `beartype_this_package`.
+    from beartype._util.func.utilfuncframe import get_frame, get_frame_package_name
+
+    # Get the calling module's name.
+    assert get_frame is not None, "get_frame is None"
+    frame = get_frame(1)
+    assert frame is not None, "frame is None"
+    calling_module_name = get_frame_package_name(frame)
+    assert calling_module_name is not None, "calling_module_name is None"
+
+    # Typecheck the calling module + any additional modules.
+    typecheck_modules((calling_module_name, *additional_modules))
+
+
+def _make_error_str(input: Any, t: Any) -> str:
+    error_components: list[str] = []
+    error_components.append("Type checking error:")
+    if hasattr(t, "__instancecheck_str__"):
+        error_components.append(t.__instancecheck_str__(input))
+    if torch is not None and torch.is_tensor(input):
+        try:
+            from lovely_tensors import lovely
+
+            error_components.append(repr(lovely(input)))
+        except BaseException:
+            error_components.append(repr(input.shape))
+    elif jax is not None and isinstance(input, jax.Array):
+        try:
+            from lovely_jax import lovely
+
+            error_components.append(repr(lovely(input)))
+        except BaseException:
+            error_components.append(repr(input.shape))
+    elif np is not None and isinstance(input, np.ndarray):
+        try:
+            from lovely_numpy import lovely
+
+            error_components.append(repr(lovely(input)))
+        except BaseException:
+            error_components.append(repr(input.shape))
+    error_components.append(shape_str(get_shape_memo()))
+
+    return "\n".join(error_components)
+
+
+T = TypeVar("T", infer_variance=True)
+
+
+def tassert(t: Any, input: T | tuple[T, ...]):
+    """
+    Typecheck the input against the given type.
+
+    Args:
+        t: Type to check against.
+        input: Input to check.
+    """
+    __tracebackhide__ = True
+
+    # Ignore typechecking if the environment variable is set.
+    if DISABLE_ENV_KEY is not None and bool(int(os.environ.get(DISABLE_ENV_KEY, "0"))):
+        return
+
+    if isinstance(input, tuple):
+        for i in input:
+            assert isinstance(i, t), _make_error_str(i, t)
+        return
+    else:
+        assert isinstance(input, t), _make_error_str(input, t)