nshutils 0.31.0__tar.gz → 0.32.0__tar.gz

{nshutils-0.31.0 → nshutils-0.32.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: nshutils
-Version: 0.31.0
+Version: 0.32.0
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshutils-0.31.0 → nshutils-0.32.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "nshutils"
-version = "0.31.0"
+version = "0.32.0"
 description = ""
 authors = [{ name = "Nima Shoghi", email = "nimashoghi@gmail.com" }]
 requires-python = ">=3.9,<4.0"

{nshutils-0.31.0 → nshutils-0.32.0}/src/nshutils/actsave/_saver.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import contextlib
 import fnmatch
+import os
 import tempfile
 import weakref
 from collections.abc import Callable, Mapping
@@ -19,7 +20,7 @@ from ..collections import apply_to_collection
 
 if not TYPE_CHECKING:
     try:
-        import torch  # type: ignore
+        import torch  # pyright: ignore[reportMissingImports]
 
         Tensor = torch.Tensor
         _torch_installed = True
@@ -29,9 +30,9 @@ if not TYPE_CHECKING:
 
         Tensor = Never
 else:
-    import torch  # type: ignore
+    import torch  # pyright: ignore[reportMissingImports]
 
-    Tensor = torch.Tensor
+    Tensor = TypeAliasType("Tensor", torch.Tensor)
     _torch_installed: Literal[True] = True
 
 log = getLogger(__name__)
@@ -59,7 +60,7 @@ def _to_numpy(activation: Value) -> np.ndarray:
         return np.array(activation)
     elif isinstance(activation, np.ndarray):
         return activation
-    elif _torch_installed and isinstance(activation, Tensor):
+    elif _torch_installed and 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
@@ -125,7 +126,8 @@ class Activation:
         if activation is None:
             return None
 
-        activation = apply_to_collection(activation, Tensor, _to_numpy)
+        if _torch_installed:
+            activation = apply_to_collection(activation, Tensor, _to_numpy)
         activation = _to_numpy(activation)
 
         # Set the transformed value
@@ -266,7 +268,11 @@ class ActSaveProvider:
 
         if save_dir is None:
             save_dir = Path(tempfile.gettempdir()) / f"actsave-{uuid7str()}"
-            log.critical(f"No save_dir specified, using {save_dir=}")
+            log.warning(
+                f"ActSave: Using temporary directory {save_dir} for activations."
+            )
+        else:
+            log.info(f"ActSave enabled. Saving to {save_dir}")
         self._saver = _Saver(save_dir, lambda: self._prefixes)
 
     def disable(self):
@@ -307,6 +313,18 @@ class ActSaveProvider:
         self._prefixes = []
         self._disable_count = 0
 
+        # Check for environment variable `ACTSAVE` to automatically enable saving.
+        # If set to "1" or "true" (case-insensitive), activations are saved to a temporary directory.
+        # If set to a path, activations are saved to that path.
+        if env_var := os.environ.get("ACTSAVE"):
+            log.info(
+                f"`ACTSAVE={env_var}` detected, attempting to auto-enable activation saving."
+            )
+            if env_var.lower() in ("1", "true"):
+                self.enable()
+            else:
+                self.enable(Path(env_var))
+
     @contextlib.contextmanager
     def disabled(self, condition: bool | Callable[[], bool] = True):
         """

{nshutils-0.31.0 → nshutils-0.32.0}/src/nshutils/lovely/_monkey_patch_all.py

@@ -30,9 +30,10 @@ def _find_deps() -> list[Library]:
 
 class monkey_patch(lovely_patch):
     def __init__(self, libraries: list[Library] | Literal["auto"] = "auto"):
-        self.libraries = libraries
-        if self.libraries == "auto":
+        if libraries == "auto":
             self.libraries = _find_deps()
+        else:
+            self.libraries = libraries
 
         if not self.libraries:
             raise ValueError(
@@ -59,7 +60,7 @@ class monkey_patch(lovely_patch):
 
                 self.stack.enter_context(numpy_monkey_patch())
             else:
-                assert_never(library)  # type: ignore
+                assert_never(library)
 
         log.info(
             f"Monkey patched libraries: {', '.join(self.libraries)}. "

{nshutils-0.31.0 → nshutils-0.32.0}/src/nshutils/lovely/jax_.py

@@ -9,7 +9,7 @@ from ._base import lovely_patch, lovely_repr
 from .utils import LovelyStats, array_stats, patch_to
 
 if TYPE_CHECKING:
-    import jax
+    import jax  # pyright: ignore[reportMissingImports]
 
 
 def _type_name(array: jax.Array):
@@ -42,7 +42,7 @@ def _dtype_str(array: jax.Array) -> str:
 
 
 def _device(array: jax.Array) -> str:
-    from jaxlib.xla_extension import Device
+    from jaxlib.xla_extension import Device  # pyright: ignore[reportMissingImports]
 
     if callable(device := array.device):
         device = device()
@@ -56,7 +56,7 @@ def _device(array: jax.Array) -> str:
 
 @lovely_repr(dependencies=["jax"])
 def jax_repr(array: jax.Array) -> LovelyStats | None:
-    import jax.numpy as jnp
+    import jax.numpy as jnp  # pyright: ignore[reportMissingImports]
 
     # For dtypes like `object` or `str`, we let the fallback repr handle it
     if not jnp.issubdtype(array.dtype, jnp.number):
@@ -85,7 +85,7 @@ class jax_monkey_patch(lovely_patch):
 
     @override
     def patch(self):
-        from jax._src import array
+        from jax._src import array  # pyright: ignore[reportMissingImports]
 
         self.prev_repr = array.ArrayImpl.__repr__
         self.prev_str = array.ArrayImpl.__str__
@@ -96,7 +96,7 @@ class jax_monkey_patch(lovely_patch):
 
     @override
     def unpatch(self):
-        from jax._src import array
+        from jax._src import array  # pyright: ignore[reportMissingImports]
 
         patch_to(array.ArrayImpl, "__repr__", self.prev_repr)
         patch_to(array.ArrayImpl, "__str__", self.prev_str)

{nshutils-0.31.0 → nshutils-0.32.0}/src/nshutils/lovely/torch_.py

@@ -9,11 +9,11 @@ from ._base import lovely_patch, lovely_repr
 from .utils import LovelyStats, array_stats, patch_to
 
 if TYPE_CHECKING:
-    import torch
+    import torch  # pyright: ignore[reportMissingImports]
 
 
 def _type_name(tensor: torch.Tensor):
-    import torch
+    import torch  # pyright: ignore[reportMissingImports]
 
     return (
         "tensor"
@@ -45,7 +45,7 @@ def _dtype_str(tensor: torch.Tensor) -> str:
 
 
 def _to_np(tensor: torch.Tensor) -> np.ndarray:
-    import torch
+    import torch  # pyright: ignore[reportMissingImports]
 
     # Get tensor data as CPU NumPy array for analysis
     t_cpu = tensor.detach().cpu()
@@ -88,7 +88,7 @@ class torch_monkey_patch(lovely_patch):
 
     @override
     def patch(self):
-        import torch
+        import torch  # pyright: ignore[reportMissingImports]
 
         self.original_repr = torch.Tensor.__repr__
         self.original_str = torch.Tensor.__str__
@@ -104,7 +104,7 @@ class torch_monkey_patch(lovely_patch):
 
     @override
     def unpatch(self):
-        import torch
+        import torch  # pyright: ignore[reportMissingImports]
 
         patch_to(torch.Tensor, "__repr__", self.original_repr)
         patch_to(torch.Tensor, "__str__", self.original_str)

{nshutils-0.31.0 → nshutils-0.32.0}/src/nshutils/snoop.py

@@ -20,19 +20,24 @@ try:
     import warnings
     from contextlib import nullcontext
 
-    import pysnooper  # type: ignore
-    import pysnooper.utils  # type: ignore
+    import pysnooper  # pyright: ignore[reportMissingImports]
+    import pysnooper.utils  # pyright: ignore[reportMissingImports]
 
     try:
-        import torch  # type: ignore
+        import torch  # pyright: ignore[reportMissingImports]
     except ImportError:
         torch = None
 
     try:
-        import numpy  # type: ignore
+        import numpy  # pyright: ignore[reportMissingImports]
     except ImportError:
         numpy = None
 
+    try:
+        import jax  # pyright: ignore[reportMissingImports]
+    except ImportError:
+        jax = None
+
     FLOATING_POINTS = set()
     for i in ["float", "double", "half", "complex128", "complex32", "complex64"]:
         # older version of PyTorch do not have complex dtypes
@@ -48,17 +53,25 @@ try:
 
     def default_format(x):
         try:
-            import lovely_tensors as lt  # type: ignore
+            from .lovely import torch_repr
 
-            return str(lt.lovely(x))
+            return torch_repr(x)
         except BaseException:
             return str(x.shape)
 
     def default_numpy_format(x):
         try:
-            import lovely_numpy as lo  # type: ignore
+            from .lovely import numpy_repr
+
+            return numpy_repr(x)
+        except BaseException:
+            return str(x.shape)
+
+    def default_jax_format(x):
+        try:
+            from .lovely import jax_repr
 
-            return str(lo.lovely(x))
+            return jax_repr(x)
         except BaseException:
             return str(x.shape)
 
@@ -68,6 +81,7 @@ try:
             *args,
             tensor_format=default_format,
             numpy_format=default_numpy_format,
+            jax_format=default_jax_format,
             **kwargs,
         ):
             self.orig_custom_repr = (
@@ -78,6 +92,7 @@ try:
             super(TorchSnooper, self).__init__(*args, **kwargs)
             self.tensor_format = tensor_format
             self.numpy_format = numpy_format
+            self.jax_format = jax_format
 
         @staticmethod
         def is_return_types(x):
@@ -176,6 +191,8 @@ try:
                 return self.tensor_format(x)
             if numpy is not None and isinstance(x, numpy.ndarray):
                 return self.numpy_format(x)
+            if jax is not None and isinstance(x, jax.Array):
+                return self.jax_format(x)
             if self.is_return_types(x):
                 return self.return_types_repr(x)
             if orig_repr_func is not repr:

{nshutils-0.31.0 → nshutils-0.32.0}/src/nshutils/typecheck.py

@@ -38,18 +38,18 @@ from jaxtyping._storage import get_shape_memo, shape_str
 from typing_extensions import TypeVar
 
 try:
-    import torch  # type: ignore
+    import torch  # pyright: ignore[reportMissingImports]
 except ImportError:
     torch = None
 
 try:
-    import np  # type: ignore
+    import np  # pyright: ignore[reportMissingImports]
 except ImportError:
     np = None
 
 
 try:
-    import jax  # type: ignore
+    import jax  # pyright: ignore[reportMissingImports]
 except ImportError:
     jax = None
 
@@ -124,23 +124,23 @@ 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  # type: ignore
+            from .lovely import torch_repr
 
-            error_components.append(repr(lovely(input)))
+            error_components.append(torch_repr(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  # type: ignore
+            from .lovely import jax_repr
 
-            error_components.append(repr(lovely(input)))
+            error_components.append(jax_repr(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  # type: ignore
+            from .lovely import numpy_repr
 
-            error_components.append(repr(lovely(input)))
+            error_components.append(numpy_repr(input))
         except BaseException:
             error_components.append(repr(input.shape))
     error_components.append(shape_str(get_shape_memo()))