nshtrainer 0.2.0__tar.gz → 0.4.0__tar.gz

{nshtrainer-0.2.0 → nshtrainer-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 0.2.0
+Version: 0.4.0
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com
@@ -9,15 +9,13 @@ 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: lightning
 Requires-Dist: lovely-numpy (>=0.2.13,<0.3.0)
 Requires-Dist: lovely-tensors (>=0.1.16,<0.2.0)
 Requires-Dist: nshconfig (>=0.2.0,<0.3.0)
 Requires-Dist: nshrunner (>=0.5.4,<0.6.0)
+Requires-Dist: nshutils (>=0.3.0,<0.4.0)
 Requires-Dist: numpy
-Requires-Dist: pysnooper
 Requires-Dist: pytorch-lightning
 Requires-Dist: rich
 Requires-Dist: torch

{nshtrainer-0.2.0 → nshtrainer-0.4.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nshtrainer"
-version = "0.2.0"
+version = "0.4.0"
 description = ""
 authors = ["Nima Shoghi <nimashoghi@gmail.com>"]
 readme = "README.md"
@@ -15,13 +15,11 @@ lightning = "*"
 pytorch-lightning = "*"
 torchmetrics = "*"
 numpy = "*"
-jaxtyping = "^0.2.33"
-beartype = "^0.18.5"
 lovely-numpy = "^0.2.13"
 lovely-tensors = "^0.1.16"
-pysnooper = "*"
 wrapt = "*"
 rich = "*"
+nshutils = "^0.3.0"
 
 
 [tool.poetry.group.dev.dependencies]

nshtrainer-0.4.0/src/nshtrainer/_snoop.py

@@ -0,0 +1 @@
+from nshutils.snoop import *  # type: ignore # noqa: F403

nshtrainer-0.4.0/src/nshtrainer/actsave/__init__.py

@@ -0,0 +1,3 @@
+from nshutils.actsave import *  # type: ignore # noqa: F403
+
+from ._callback import ActSaveCallback as ActSaveCallback

{nshtrainer-0.2.0 → nshtrainer-0.4.0}/src/nshtrainer/actsave/_callback.py

@@ -3,10 +3,9 @@ from typing import TYPE_CHECKING, Literal, cast
 
 from lightning.pytorch import LightningModule, Trainer
 from lightning.pytorch.callbacks.callback import Callback
+from nshutils.actsave import ActSave
 from typing_extensions import TypeAlias, override
 
-from ._saver import ActSave
-
 if TYPE_CHECKING:
     from ..model.config import BaseConfig
 

nshtrainer-0.4.0/src/nshtrainer/typecheck.py

@@ -0,0 +1 @@
+from nshutils.typecheck import *  # type: ignore # noqa: F403

nshtrainer-0.2.0/src/nshtrainer/_snoop.py

@@ -1,216 +0,0 @@
-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)

nshtrainer-0.2.0/src/nshtrainer/actsave/__init__.py

@@ -1,7 +0,0 @@
-from ._callback import ActSaveCallback as ActSaveCallback
-from ._loader import ActivationLoader as ActivationLoader
-from ._loader import ActLoad as ActLoad
-from ._saver import Activation as Activation
-from ._saver import ActivationSaver as ActivationSaver
-from ._saver import ActSave as ActSave
-from ._saver import Transform as Transform

nshtrainer-0.2.0/src/nshtrainer/actsave/_loader.py

@@ -1,144 +0,0 @@
-import pprint
-from dataclasses import dataclass, field
-from functools import cached_property
-from logging import getLogger
-from pathlib import Path
-from typing import cast, overload
-
-import numpy as np
-from typing_extensions import TypeVar, override
-
-log = getLogger(__name__)
-
-T = TypeVar("T", infer_variance=True)
-
-
-@dataclass
-class LoadedActivation:
-    base_dir: Path = field(repr=False)
-    name: str
-    num_activations: int = field(init=False)
-    activation_files: list[Path] = field(init=False, repr=False)
-
-    def __post_init__(self):
-        if not self.activation_dir.exists():
-            raise ValueError(f"Activation dir {self.activation_dir} does not exist")
-
-        # The number of activations = the * of .npy files in the activation dir
-        self.activation_files = list(self.activation_dir.glob("*.npy"))
-        # Sort the activation files by the numerical index in the filename
-        self.activation_files.sort(key=lambda p: int(p.stem))
-        self.num_activations = len(self.activation_files)
-
-    @property
-    def activation_dir(self) -> Path:
-        return self.base_dir / self.name
-
-    def _load_activation(self, item: int):
-        activation_path = self.activation_files[item]
-        if not activation_path.exists():
-            raise ValueError(f"Activation {activation_path} does not exist")
-        return cast(np.ndarray, np.load(activation_path, allow_pickle=True))
-
-    @overload
-    def __getitem__(self, item: int) -> np.ndarray: ...
-
-    @overload
-    def __getitem__(self, item: slice | list[int]) -> list[np.ndarray]: ...
-
-    def __getitem__(
-        self, item: int | slice | list[int]
-    ) -> np.ndarray | list[np.ndarray]:
-        if isinstance(item, int):
-            return self._load_activation(item)
-        elif isinstance(item, slice):
-            return [
-                self._load_activation(i)
-                for i in range(*item.indices(self.num_activations))
-            ]
-        elif isinstance(item, list):
-            return [self._load_activation(i) for i in item]
-        else:
-            raise TypeError(f"Invalid type {type(item)} for item {item}")
-
-    def __iter__(self):
-        return iter(self[i] for i in range(self.num_activations))
-
-    def __len__(self):
-        return self.num_activations
-
-    def all_activations(self):
-        return [self[i] for i in range(self.num_activations)]
-
-    @override
-    def __repr__(self):
-        return f"<LoadedActivation {self.name} ({self.num_activations} activations)>"
-
-
-class ActLoad:
-    @classmethod
-    def all_versions(cls, dir: str | Path):
-        dir = Path(dir)
-
-        # If the dir is not an activation base directory, we return None
-        if not (dir / ".activationbase").exists():
-            return None
-
-        # The contents of `dir` should be directories, each of which is a version.
-        return [
-            (subdir, int(subdir.name)) for subdir in dir.iterdir() if subdir.is_dir()
-        ]
-
-    @classmethod
-    def is_valid_activation_base(cls, dir: str | Path):
-        return cls.all_versions(dir) is not None
-
-    @classmethod
-    def from_latest_version(cls, dir: str | Path):
-        # The contents of `dir` should be directories, each of which is a version
-        # We need to find the latest version
-        if (all_versions := cls.all_versions(dir)) is None:
-            raise ValueError(f"{dir} is not an activation base directory")
-
-        path, _ = max(all_versions, key=lambda p: p[1])
-        return cls(path)
-
-    def __init__(self, dir: Path):
-        self._dir = dir
-
-    def activation(self, name: str):
-        return LoadedActivation(self._dir, name)
-
-    @cached_property
-    def activations(self):
-        dirs = list(self._dir.iterdir())
-        # Sort the dirs by the last modified time
-        dirs.sort(key=lambda p: p.stat().st_mtime)
-
-        return {p.name: LoadedActivation(self._dir, p.name) for p in dirs}
-
-    def __iter__(self):
-        return iter(self.activations.values())
-
-    def __getitem__(self, item: str):
-        return self.activations[item]
-
-    def __len__(self):
-        return len(self.activations)
-
-    @override
-    def __repr__(self):
-        acts_str = pprint.pformat(
-            {
-                name: f"<{activation.num_activations} activations>"
-                for name, activation in self.activations.items()
-            }
-        )
-        acts_str = acts_str.replace("'<", "<").replace(">'", ">")
-        return f"ActLoad({acts_str})"
-
-    def get(self, name: str, /, default: T) -> LoadedActivation | T:
-        return self.activations.get(name, default)
-
-
-ActivationLoader = ActLoad

nshtrainer-0.2.0/src/nshtrainer/actsave/_saver.py

@@ -1,337 +0,0 @@
-import contextlib
-import fnmatch
-import tempfile
-import uuid
-import weakref
-from collections.abc import Callable, Mapping
-from dataclasses import dataclass
-from functools import wraps
-from logging import getLogger
-from pathlib import Path
-from typing import Generic, TypeAlias, cast, overload
-
-import numpy as np
-import torch
-from lightning_utilities.core.apply_func import apply_to_collection
-from typing_extensions import ParamSpec, TypeVar, override
-
-log = getLogger(__name__)
-
-Value: TypeAlias = int | float | complex | bool | str | np.ndarray | torch.Tensor | None
-ValueOrLambda = Value | Callable[..., Value]
-
-
-def _to_numpy(activation: Value) -> np.ndarray:
-    # Make sure it's not `None`
-    if activation is None:
-        raise ValueError("Activation should not be `None`")
-
-    if isinstance(activation, np.ndarray):
-        return activation
-    if isinstance(activation, torch.Tensor):
-        activation = activation.detach()
-        if activation.is_floating_point():
-            # NOTE: We need to convert to float32 because [b]float16 is not supported by numpy
-            activation = activation.float()
-        return activation.cpu().numpy()
-    if isinstance(activation, (int, float, complex, str, bool)):
-        return np.array(activation)
-
-    return activation
-
-
-T = TypeVar("T", infer_variance=True)
-
-
-# A wrapper around weakref.ref that allows for primitive types
-# To get around errors like:
-# TypeError: cannot create weak reference to 'int' object
-class WeakRef(Generic[T]):
-    _ref: Callable[[], T] | None
-
-    def __init__(self, obj: T):
-        try:
-            self._ref = cast(Callable[[], T], weakref.ref(obj))
-        except TypeError as e:
-            if "cannot create weak reference" not in str(e):
-                raise
-            self._ref = lambda: obj
-
-    def __call__(self) -> T:
-        if self._ref is None:
-            raise RuntimeError("WeakRef is deleted")
-        return self._ref()
-
-    def delete(self):
-        del self._ref
-        self._ref = None
-
-
-@dataclass
-class Activation:
-    name: str
-    ref: WeakRef[ValueOrLambda] | None
-    transformed: np.ndarray | None = None
-
-    def __post_init__(self):
-        # Update the `name` to replace `/` with `.`
-        self.name = self.name.replace("/", ".")
-
-    def __call__(self) -> np.ndarray | None:
-        # If we have a transformed value, we return it
-        if self.transformed is not None:
-            return self.transformed
-
-        if self.ref is None:
-            raise RuntimeError("Activation is deleted")
-
-        # If we have a lambda, we need to call it
-        unrwapped_ref = self.ref()
-        activation = unrwapped_ref
-        if callable(unrwapped_ref):
-            activation = unrwapped_ref()
-
-        # If we have a `None`, we return early
-        if activation is None:
-            return None
-
-        activation = apply_to_collection(activation, torch.Tensor, _to_numpy)
-        activation = _to_numpy(activation)
-
-        # Set the transformed value
-        self.transformed = activation
-
-        # Delete the reference
-        self.ref.delete()
-        del self.ref
-        self.ref = None
-
-        return self.transformed
-
-    @classmethod
-    def from_value_or_lambda(cls, name: str, value_or_lambda: ValueOrLambda):
-        return cls(name, WeakRef(value_or_lambda))
-
-    @classmethod
-    def from_dict(cls, d: Mapping[str, ValueOrLambda]):
-        return [cls.from_value_or_lambda(k, v) for k, v in d.items()]
-
-
-Transform = Callable[[Activation], Mapping[str, ValueOrLambda]]
-
-
-def _ensure_supported():
-    try:
-        import torch.distributed as dist
-
-        if dist.is_initialized() and dist.get_world_size() > 1:
-            raise RuntimeError("Only single GPU is supported at the moment")
-    except ImportError:
-        pass
-
-
-P = ParamSpec("P")
-
-
-def _ignore_if_scripting(fn: Callable[P, None]) -> Callable[P, None]:
-    @wraps(fn)
-    def wrapper(*args: P.args, **kwargs: P.kwargs) -> None:
-        if torch.jit.is_scripting():
-            return
-
-        _ensure_supported()
-        fn(*args, **kwargs)
-
-    return wrapper
-
-
-class _Saver:
-    def __init__(
-        self,
-        save_dir: Path,
-        prefixes_fn: Callable[[], list[str]],
-        *,
-        filters: list[str] | None = None,
-    ):
-        # Create a directory under `save_dir` by autoincrementing
-        # (i.e., every activation save context, we create a new directory)
-        # The id = the number of activation subdirectories
-        self._id = sum(1 for subdir in save_dir.glob("*") if subdir.is_dir())
-        save_dir.mkdir(parents=True, exist_ok=True)
-
-        # Add a .activationbase file to the save_dir to indicate that this is an activation base
-        (save_dir / ".activationbase").touch(exist_ok=True)
-
-        self._save_dir = save_dir / f"{self._id:04d}"
-        # Make sure `self._save_dir` does not exist and create it
-        self._save_dir.mkdir(exist_ok=False)
-
-        self._prefixes_fn = prefixes_fn
-        self._filters = filters
-
-    def _save_activation(self, activation: Activation):
-        # If the activation value is `None`, we skip it.
-        if (activation_value := activation()) is None:
-            return
-
-        # Save the activation to self._save_dir / name / {id}.npz, where id is an auto-incrementing integer
-        file_name = ".".join(self._prefixes_fn() + [activation.name])
-        path = self._save_dir / file_name
-        path.mkdir(exist_ok=True, parents=True)
-
-        # Get the next id and save the activation
-        id = len(list(path.glob("*.npy")))
-        np.save(path / f"{id:04d}.npy", activation_value)
-
-    @_ignore_if_scripting
-    def save(
-        self,
-        acts: dict[str, ValueOrLambda] | None = None,
-        /,
-        **kwargs: ValueOrLambda,
-    ):
-        kwargs.update(acts or {})
-
-        # Build activations
-        activations = Activation.from_dict(kwargs)
-
-        for activation in activations:
-            # Make sure name matches at least one filter if filters are specified
-            if self._filters is not None and all(
-                not fnmatch.fnmatch(activation.name, f) for f in self._filters
-            ):
-                continue
-
-            # Save the current activation
-            self._save_activation(activation)
-
-        del activations
-
-
-class ActSaveProvider:
-    _saver: _Saver | None = None
-    _prefixes: list[str] = []
-
-    def initialize(self, save_dir: Path | None = None):
-        """
-        Initializes the saver with the given configuration and save directory.
-
-        Args:
-            save_dir (Path): The directory where the saved files will be stored.
-        """
-        if self._saver is None:
-            if save_dir is None:
-                save_dir = Path(tempfile.gettempdir()) / f"actsave-{uuid.uuid4()}"
-                log.critical(f"No save_dir specified, using {save_dir=}")
-            self._saver = _Saver(
-                save_dir,
-                lambda: self._prefixes,
-            )
-
-    @contextlib.contextmanager
-    def enabled(self, save_dir: Path | None = None):
-        """
-        Context manager that enables the actsave functionality with the specified configuration.
-
-        Args:
-            save_dir (Path): The directory where the saved files will be stored.
-        """
-        prev = self._saver
-        self.initialize(save_dir)
-        try:
-            yield
-        finally:
-            self._saver = prev
-
-    @override
-    def __init__(self):
-        super().__init__()
-
-        self._saver = None
-        self._prefixes = []
-
-    @contextlib.contextmanager
-    def context(self, label: str):
-        """
-        A context manager that adds a label to the current context.
-
-        Args:
-            label (str): The label for the context.
-        """
-        if torch.jit.is_scripting():
-            yield
-            return
-
-        if self._saver is None:
-            yield
-            return
-
-        _ensure_supported()
-
-        log.debug(f"Entering ActSave context {label}")
-        self._prefixes.append(label)
-        try:
-            yield
-        finally:
-            _ = self._prefixes.pop()
-
-    prefix = context
-
-    @overload
-    def __call__(
-        self,
-        acts: dict[str, ValueOrLambda] | None = None,
-        /,
-        **kwargs: ValueOrLambda,
-    ):
-        """
-        Saves the activations to disk.
-
-        Args:
-            acts (dict[str, ValueOrLambda] | None, optional): A dictionary of acts. Defaults to None.
-            **kwargs (ValueOrLambda): Additional keyword arguments.
-
-        Returns:
-            None
-
-        """
-        ...
-
-    @overload
-    def __call__(self, acts: Callable[[], dict[str, ValueOrLambda]], /):
-        """
-        Saves the activations to disk.
-
-        Args:
-            acts (Callable[[], dict[str, ValueOrLambda]]): A callable that returns a dictionary of acts.
-            **kwargs (ValueOrLambda): Additional keyword arguments.
-
-        Returns:
-            None
-
-        """
-        ...
-
-    def __call__(
-        self,
-        acts: (
-            dict[str, ValueOrLambda] | Callable[[], dict[str, ValueOrLambda]] | None
-        ) = None,
-        /,
-        **kwargs: ValueOrLambda,
-    ):
-        if torch.jit.is_scripting():
-            return
-
-        if self._saver is None:
-            return
-
-        if acts is not None and callable(acts):
-            acts = acts()
-        self._saver.save(acts, **kwargs)
-
-    save = __call__
-
-
-ActSave = ActSaveProvider()
-ActivationSaver = ActSave

nshtrainer-0.2.0/src/nshtrainer/typecheck.py

@@ -1,145 +0,0 @@
-import os
-from collections.abc import Sequence
-from logging import getLogger
-from typing import Any
-
-import numpy as np
-import torch
-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 torch import Tensor as Tensor
-from torch.nn.parameter import Parameter as Parameter
-from typing_extensions import TypeVar
-
-log = getLogger(__name__)
-
-DISABLE_ENV_KEY = "LL_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_tensor(input):
-        try:
-            from lovely_tensors 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", torch.Tensor, np.ndarray, infer_variance=True)
-
-"""
-Patch to jaxtyping:
-
-In `jaxtyping._import_hook`, we add:
-def _has_isinstance_or_tassert(func_def):
-    for node in ast.walk(func_def):
-        if isinstance(node, ast.Call):
-            if isinstance(node.func, ast.Name) and node.func.id == "isinstance":
-                return True
-            elif isinstance(node.func, ast.Name) and node.func.id == "tassert":
-                return True
-    return False
-
-and we check this when adding the decorators.
-"""
-
-
-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.
-    """
-
-    # 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)