nshutils 0.5.1__tar.gz → 0.6.1__tar.gz

{nshutils-0.5.1 → nshutils-0.6.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshutils
-Version: 0.5.1
+Version: 0.6.1
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshutils-0.5.1 → nshutils-0.6.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nshutils"
-version = "0.5.1"
+version = "0.6.1"
 description = ""
 authors = ["Nima Shoghi <nimashoghi@gmail.com>"]
 readme = "README.md"
@@ -14,6 +14,12 @@ beartype = "^0.18.5"
 numpy = "*"
 
 
+[tool.poetry.group.dev.dependencies]
+pyright = "^1.1.373"
+ruff = "^0.5.4"
+ipykernel = "^6.29.5"
+ipywidgets = "^8.1.3"
+
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"

{nshutils-0.5.1 → nshutils-0.6.1}/src/nshutils/actsave/_saver.py

@@ -16,7 +16,7 @@ from typing_extensions import Never, ParamSpec, TypeVar, override
 from ..collections import apply_to_collection
 
 try:
-    import torch
+    import torch  # type: ignore
 
     if not TYPE_CHECKING:
         Tensor: TypeAlias = torch.Tensor
@@ -145,7 +145,7 @@ Transform = Callable[[Activation], Mapping[str, ValueOrLambda]]
 
 def _ensure_supported():
     try:
-        import torch.distributed as dist
+        import torch.distributed as dist  # type: ignore
 
         if dist.is_initialized() and dist.get_world_size() > 1:
             raise RuntimeError("Only single GPU is supported at the moment")

{nshutils-0.5.1 → nshutils-0.6.1}/src/nshutils/collections.py

@@ -6,7 +6,7 @@ import dataclasses
 from collections import OrderedDict, defaultdict
 from collections.abc import Callable, Mapping, Sequence
 from copy import deepcopy
-from typing import Any
+from typing import Any, cast
 
 
 def is_namedtuple(obj: object) -> bool:
@@ -97,6 +97,7 @@ def apply_to_collection(
         return elem_type(*out) if is_namedtuple_ else elem_type(out)
 
     if is_dataclass_instance(data):
+        data = cast(Any, data)
         # make a deepcopy of the data,
         # but do not deepcopy mapped fields since the computation would
         # be wasted on values that likely get immediately overwritten
@@ -136,136 +137,3 @@ def apply_to_collection(
 
     # data is neither of dtype, nor a collection
     return data
-
-
-def apply_to_collections(
-    data1: Any | None,
-    data2: Any | None,
-    dtype: type | Any | tuple[type | Any],
-    function: Callable,
-    *args: Any,
-    wrong_dtype: type | tuple[type] | None = None,
-    **kwargs: Any,
-) -> Any:
-    """Zips two collections and applies a function to their items of a certain dtype.
-
-    Args:
-        data1: The first collection
-        data2: The second collection
-        dtype: the given function will be applied to all elements of this dtype
-        function: the function to apply
-        *args: positional arguments (will be forwarded to calls of ``function``)
-        wrong_dtype: the given function won't be applied if this type is specified and the given collections
-            is of the ``wrong_dtype`` even if it is of type ``dtype``
-        **kwargs: keyword arguments (will be forwarded to calls of ``function``)
-
-    Returns:
-        The resulting collection
-
-    Raises:
-        AssertionError:
-            If sequence collections have different data sizes.
-    """
-    if data1 is None:
-        if data2 is None:
-            return None
-        # in case they were passed reversed
-        data1, data2 = data2, None
-
-    elem_type = type(data1)
-
-    if (
-        isinstance(data1, dtype)
-        and data2 is not None
-        and (wrong_dtype is None or not isinstance(data1, wrong_dtype))
-    ):
-        return function(data1, data2, *args, **kwargs)
-
-    if isinstance(data1, Mapping) and data2 is not None:
-        # use union because we want to fail if a key does not exist in both
-        zipped = {k: (data1[k], data2[k]) for k in data1.keys() | data2.keys()}
-        return elem_type(
-            {
-                k: apply_to_collections(
-                    *v, dtype, function, *args, wrong_dtype=wrong_dtype, **kwargs
-                )
-                for k, v in zipped.items()
-            }
-        )
-
-    is_namedtuple_ = is_namedtuple(data1)
-    is_sequence = isinstance(data1, Sequence) and not isinstance(data1, str)
-    if (is_namedtuple_ or is_sequence) and data2 is not None:
-        if len(data1) != len(data2):
-            raise ValueError("Sequence collections have different sizes.")
-        out = [
-            apply_to_collections(
-                v1, v2, dtype, function, *args, wrong_dtype=wrong_dtype, **kwargs
-            )
-            for v1, v2 in zip(data1, data2)
-        ]
-        return elem_type(*out) if is_namedtuple_ else elem_type(out)
-
-    if is_dataclass_instance(data1) and data2 is not None:
-        if not is_dataclass_instance(data2):
-            raise TypeError(
-                "Expected inputs to be dataclasses of the same type or to have identical fields"
-                f" but got input 1 of type {type(data1)} and input 2 of type {type(data2)}."
-            )
-        if not (
-            len(dataclasses.fields(data1)) == len(dataclasses.fields(data2))
-            and all(
-                map(
-                    lambda f1, f2: isinstance(f1, type(f2)),
-                    dataclasses.fields(data1),
-                    dataclasses.fields(data2),
-                )
-            )
-        ):
-            raise TypeError("Dataclasses fields do not match.")
-        # make a deepcopy of the data,
-        # but do not deepcopy mapped fields since the computation would
-        # be wasted on values that likely get immediately overwritten
-        data = [data1, data2]
-        fields: list[dict] = [{}, {}]
-        memo: dict = {}
-        for i in range(len(data)):
-            for field in dataclasses.fields(data[i]):
-                field_value = getattr(data[i], field.name)
-                fields[i][field.name] = (field_value, field.init)
-                if i == 0:
-                    memo[id(field_value)] = field_value
-
-        result = deepcopy(data1, memo=memo)
-
-        # apply function to each field
-        for (field_name, (field_value1, field_init1)), (
-            _,
-            (field_value2, field_init2),
-        ) in zip(fields[0].items(), fields[1].items()):
-            v = None
-            if field_init1 and field_init2:
-                v = apply_to_collections(
-                    field_value1,
-                    field_value2,
-                    dtype,
-                    function,
-                    *args,
-                    wrong_dtype=wrong_dtype,
-                    **kwargs,
-                )
-            if not field_init1 or not field_init2 or v is None:  # retain old value
-                return apply_to_collection(
-                    data1, dtype, function, *args, wrong_dtype=wrong_dtype, **kwargs
-                )
-            try:
-                setattr(result, field_name, v)
-            except dataclasses.FrozenInstanceError as e:
-                raise ValueError(
-                    "A frozen dataclass was passed to `apply_to_collections` but this is not allowed."
-                ) from e
-        return result
-
-    return apply_to_collection(
-        data1, dtype, function, *args, wrong_dtype=wrong_dtype, **kwargs
-    )

{nshutils-0.5.1 → nshutils-0.6.1}/src/nshutils/logging.py

@@ -13,7 +13,7 @@ def init_python_logging(
 ):
     if lovely_tensors:
         try:
-            import lovely_tensors as _lovely_tensors
+            import lovely_tensors as _lovely_tensors  # type: ignore
 
             _lovely_tensors.monkey_patch()
         except ImportError:
@@ -23,7 +23,7 @@ def init_python_logging(
 
     if lovely_numpy:
         try:
-            import lovely_numpy as _lovely_numpy
+            import lovely_numpy as _lovely_numpy  # type: ignore
 
             _lovely_numpy.set_config(repr=_lovely_numpy.lovely)
         except ImportError:
@@ -39,7 +39,7 @@ def init_python_logging(
 
     if rich:
         try:
-            from rich.logging import RichHandler
+            from rich.logging import RichHandler  # type: ignore
 
             log_handlers.append(RichHandler(rich_tracebacks=rich_tracebacks))
         except ImportError:

{nshutils-0.5.1 → nshutils-0.6.1}/src/nshutils/snoop.py

@@ -21,19 +21,19 @@ try:
     from pkg_resources import DistributionNotFound, get_distribution
 
     try:
-        import torch
+        import torch  # type: ignore
     except ImportError:
         torch = None
 
     try:
-        import numpy
+        import numpy  # type: ignore
     except ImportError:
         numpy = None
 
     FLOATING_POINTS = set()
     for i in ["float", "double", "half", "complex128", "complex32", "complex64"]:
         # older version of PyTorch do not have complex dtypes
-        if torch is not None and not hasattr(torch, i):
+        if torch is None or not hasattr(torch, i):
             continue
         FLOATING_POINTS.add(getattr(torch, i))
 
@@ -45,15 +45,15 @@ try:
 
     def default_format(x):
         try:
-            import lovely_tensors as lt
+            import lovely_tensors as lt  # type: ignore
 
-            return = str(lt.lovely(x))
+            return str(lt.lovely(x))
         except BaseException:
             return str(x.shape)
 
     def default_numpy_format(x):
         try:
-            import lovely_numpy as lo
+            import lovely_numpy as lo  # type: ignore
 
             return str(lo.lovely(x))
         except BaseException:

{nshutils-0.5.1 → nshutils-0.6.1}/src/nshutils/typecheck.py

@@ -34,18 +34,18 @@ from jaxtyping._storage import get_shape_memo, shape_str
 from typing_extensions import TypeVar
 
 try:
-    import torch
+    import torch  # type: ignore
 except ImportError:
     torch = None
 
 try:
-    import np
+    import np  # type: ignore
 except ImportError:
     np = None
 
 
 try:
-    import jax
+    import jax  # type: ignore
 except ImportError:
     jax = None
 log = getLogger(__name__)
@@ -106,21 +106,21 @@ def _make_error_str(input: Any, t: Any) -> str:
         error_components.append(t.__instancecheck_str__(input))
     if torch is not None and torch.is_tensor(input):
         try:
-            from lovely_tensors import lovely
+            from lovely_tensors import lovely  # type: ignore
 
             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
+            from lovely_jax import lovely  # type: ignore
 
             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
+            from lovely_numpy import lovely  # type: ignore
 
             error_components.append(repr(lovely(input)))
         except BaseException: