nshtrainer 0.1.1__py3-none-any.whl → 0.3.0__py3-none-any.whl

nshtrainer/__init__.py

@@ -1,11 +1,11 @@
 from . import _experimental as _experimental
 from . import actsave as actsave
 from . import callbacks as callbacks
+from . import config as config
 from . import lr_scheduler as lr_scheduler
 from . import model as model
 from . import nn as nn
 from . import optimizer as optimizer
-from . import snapshot as snapshot
 from . import typecheck as typecheck
 from ._snoop import snoop as snoop
 from .actsave import ActLoad as ActLoad

nshtrainer/_snoop.py

@@ -1,216 +1 @@
-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)
+from nshutils.snoop import *  # type: ignore # noqa: F403

nshtrainer/config.py

@@ -0,0 +1,4 @@
+from nshconfig import *  # type: ignore # noqa: F403
+from nshconfig import Config
+
+TypedConfig = Config

nshtrainer/runner.py

@@ -7,13 +7,13 @@ from typing_extensions import TypeVar, TypeVarTuple, Unpack, override
 from .model.config import BaseConfig
 
 TConfig = TypeVar("TConfig", bound=BaseConfig, infer_variance=True)
-TArguments = TypeVarTuple("TArguments")
+TArguments = TypeVarTuple("TArguments", default=Unpack[tuple[()]])
 TReturn = TypeVar("TReturn", infer_variance=True)
 
 
 class Runner(
-    _Runner[Unpack[tuple[TConfig, Unpack[TArguments]]], TReturn],
-    Generic[TConfig, Unpack[TArguments], TReturn],
+    _Runner[TReturn, TConfig, Unpack[TArguments]],
+    Generic[TReturn, TConfig, Unpack[TArguments]],
 ):
     @override
     @classmethod

nshtrainer/typecheck.py

@@ -1,145 +1 @@
-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)
+from nshutils.typecheck import *  # type: ignore # noqa: F403

{nshtrainer-0.1.1.dist-info → nshtrainer-0.3.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 0.1.1
+Version: 0.3.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.3,<0.6.0)
+Requires-Dist: nshrunner (>=0.5.4,<0.6.0)
+Requires-Dist: nshutils (>=0.2.0,<0.3.0)
 Requires-Dist: numpy
-Requires-Dist: pysnooper
 Requires-Dist: pytorch-lightning
 Requires-Dist: rich
 Requires-Dist: torch

{nshtrainer-0.1.1.dist-info → nshtrainer-0.3.0.dist-info}/RECORD

@@ -1,9 +1,9 @@
-nshtrainer/__init__.py,sha256=OHbxLxVvFGW--ecuIGqkoylSVHFS4x4F1-oeuENH-Do,2212
+nshtrainer/__init__.py,sha256=_r7kBmgGSLVfActlqQeupNolrmBu45xUuSS8odt3HL8,2208
 nshtrainer/_experimental/__init__.py,sha256=2tQIcrWT8U8no_AeBTYnozaTmxN40kuAJdGQ4b-PoWM,120
 nshtrainer/_experimental/flops/__init__.py,sha256=edo9Ez3LlrnxkNRX9W6YBhPkRPKYGLpkpnl5gx7sEX8,1550
 nshtrainer/_experimental/flops/flop_counter.py,sha256=-sL0Fy6poXa__hyzUMdZScjPULp4coQELQpPU6p6dXU,25736
 nshtrainer/_experimental/flops/module_tracker.py,sha256=bUL-IRTd0aF_DwmXkZjHZAA31p4ZEhyqhc26XWKQUUY,4922
-nshtrainer/_snoop.py,sha256=Rofv1Rd92E0LY40G3A-o9Hu0ZI73RR59wJD5l4Q3PDM,7022
+nshtrainer/_snoop.py,sha256=2rEemPyMP3aIo2QgPzo_-AlT1oXGWYQipId4RQskMls,58
 nshtrainer/actsave/__init__.py,sha256=G1T-fELuGWkVqdhdyoePtj2dTOUtcIOW4VgsXv9JNTA,338
 nshtrainer/actsave/_callback.py,sha256=QoTa60F70f1RxB41VKixN9l5_htfFQxXDPHHSNFreuk,2770
 nshtrainer/actsave/_loader.py,sha256=fAhD32DrJa4onkYfcwc21YIeGEYzOSXCK_HVo9SZLgQ,4604
@@ -24,6 +24,7 @@ nshtrainer/callbacks/print_table.py,sha256=FcA-CBWwMf9c1NNRinvYpZC400RNQxuP28bJf
 nshtrainer/callbacks/throughput_monitor.py,sha256=YQLdpX3LGybIiD814yT9yCCVSEXRWf8WwsvVaN5aDBE,1848
 nshtrainer/callbacks/timer.py,sha256=sDXPPcdDKu5xnuK_bjr8plIq9MBuluNJ42Mt9LvPZzc,4610
 nshtrainer/callbacks/wandb_watch.py,sha256=pUpMsNxd03ex1rzOmFw2HzGOXjnQGaH84m8cc2dXo4g,2937
+nshtrainer/config.py,sha256=IXOAl_JWFNX9kPTo_iw4Nc3qXqkKrbA6-ZrvTAjqu6A,104
 nshtrainer/data/__init__.py,sha256=7mk1tr7SWUZ7ySbsf0y0ZPszk7u4QznPhQ-7wnpH9ec,149
 nshtrainer/data/balanced_batch_sampler.py,sha256=bcJBcQjh1hB1yKF_xSlT9AtEWv0BJjYc1CuH2BF-ea8,4392
 nshtrainer/data/transform.py,sha256=JeGxvytQly8hougrsdMmKG8gJ6qvFPDglJCO4Tp6STk,1795
@@ -47,18 +48,18 @@ nshtrainer/nn/module_dict.py,sha256=NOY0B6WDTnktyWH4GthsprMQo0bpehC-hCq9SfD8paE,
 nshtrainer/nn/module_list.py,sha256=fb2u5Rqdjff8Pekyr9hkCPkBorQ-fldzzFAjsgWAm30,1719
 nshtrainer/nn/nonlinearity.py,sha256=owtU4kh4G98psD0axOJWVfBhm-OtJVgFM-TXSHmbNPU,3625
 nshtrainer/optimizer.py,sha256=kuJEA1pvB3y1FcsfhAoOJujVqEZqFHlmYO8GW6JeA1g,1527
-nshtrainer/runner.py,sha256=af_EGnQTSvUgwnVhhytvY3V7o_Xg-xx-sLb8K2Szb1E,979
+nshtrainer/runner.py,sha256=vyHr0EZ0PBOWZh09BtOOxio-FRQZFbVoL4cdBlI97vY,991
 nshtrainer/scripts/check_env.py,sha256=IMl6dSqsLYppI0XuCsVq8lK4bYqXwY9KHJkzsShz4Kg,806
 nshtrainer/scripts/find_packages.py,sha256=FbdlfmAefttFSMfaT0A46a-oHLP_ioaQKihwBfBeWeA,1467
 nshtrainer/trainer/__init__.py,sha256=P2rmr8oBVTHk-HJHYPcUwWqDEArMbPR4_rPpATbWK3E,40
 nshtrainer/trainer/signal_connector.py,sha256=QAoPM_C5JJOVQebcrJOimUUD3GHyoeZUqCEAvzZlT4U,8710
 nshtrainer/trainer/trainer.py,sha256=eYEYfY9v70MuorHcSf8nqM7f2CkmUHhpPcjCk4FJD7k,14034
-nshtrainer/typecheck.py,sha256=RGYHxDBcs97E6ayl6Olc43JBZXQolCtMxcLBniVCVBg,4688
+nshtrainer/typecheck.py,sha256=ryV1Tzcf7hJ4I19H1oQVkikU9spmRk8jyIKQZ5UF7pQ,62
 nshtrainer/util/environment.py,sha256=_SEtiQ_s5bL5pllUlf96AOUv15kNvCPvocVC13S7mIk,4166
 nshtrainer/util/seed.py,sha256=HEXgVs-wldByahOysKwq7506OHxdYTEgmP-tDQVAEkQ,287
 nshtrainer/util/slurm.py,sha256=rofIU26z3SdL79SF45tNez6juou1cyDLz07oXEZb9Hg,1566
 nshtrainer/util/typed.py,sha256=NGuDkDzFlc1fAoaXjOFZVbmj0mRFjsQi1E_hPa7Bn5U,128
 nshtrainer/util/typing_utils.py,sha256=8ptjSSLZxlmy4FY6lzzkoGoF5fGNClo8-B_c0XHQaNU,385
-nshtrainer-0.1.1.dist-info/METADATA,sha256=32iVLvdJh6OJQyD-_7NDO6IYqfHPSflDznYfYaCo8-c,882
-nshtrainer-0.1.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.1.1.dist-info/RECORD,,
+nshtrainer-0.3.0.dist-info/METADATA,sha256=oME_P_Y7U4bavZunb-rF2m1R1w31vOS0Qh9l0Nfua68,812
+nshtrainer-0.3.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.3.0.dist-info/RECORD,,