nshutils 0.22.5__tar.gz → 0.30.1__tar.gz

{nshutils-0.22.5 → nshutils-0.30.1}/PKG-INFO

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

{nshutils-0.22.5 → nshutils-0.30.1}/pyproject.toml

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

{nshutils-0.22.5 → nshutils-0.30.1}/src/nshutils/actsave/_saver.py

@@ -241,6 +241,17 @@ class _Saver:
 class ActSaveProvider:
     _saver: _Saver | None = None
     _prefixes: list[str] = []
+    _disable_count: int = 0
+
+    @property
+    def is_initialized(self) -> bool:
+        """Returns True if ActSave.enable() has been called and not subsequently disabled."""
+        return self._saver is not None
+
+    @property
+    def is_enabled(self) -> bool:
+        """Returns True if ActSave is currently active and will save activations."""
+        return self.is_initialized and self._disable_count == 0
 
     def enable(self, save_dir: Path | None = None):
         """
@@ -294,6 +305,34 @@ class ActSaveProvider:
 
         self._saver = None
         self._prefixes = []
+        self._disable_count = 0
+
+    @contextlib.contextmanager
+    def disabled(self, condition: bool | Callable[[], bool] = True):
+        """
+        Context manager to temporarily disable activation saving.
+
+        Args:
+            condition (bool | Callable[[], bool], optional):
+                If True or a callable returning True, saving is disabled within this context.
+                Defaults to True.
+        """
+        if _torch_is_scripting():
+            yield
+            return
+
+        should_disable = condition() if callable(condition) else condition
+        if should_disable:
+            self._disable_count += 1
+
+        try:
+            yield
+        finally:
+            if should_disable:
+                self._disable_count -= 1
+                if self._disable_count < 0:  # Should not happen
+                    log.warning("ActSave disable count went below zero.")
+                    self._disable_count = 0
 
     @contextlib.contextmanager
     def context(self, label: str):
@@ -307,7 +346,7 @@ class ActSaveProvider:
             yield
             return
 
-        if self._saver is None:
+        if not self.is_enabled:
             yield
             return
 
@@ -368,9 +407,12 @@ class ActSaveProvider:
         if _torch_is_scripting():
             return
 
-        if self._saver is None:
+        if not self.is_enabled:
             return
 
+        # Ensure _saver is not None, which is guaranteed by is_enabled but mypy needs help
+        assert self._saver is not None
+
         if acts is not None and callable(acts):
             acts = acts()
         self._saver.save(acts, **kwargs)

{nshutils-0.22.5 → nshutils-0.30.1}/src/nshutils/lovely/_base.py

@@ -4,8 +4,16 @@ import functools
 import importlib.util
 import logging
 from collections.abc import Callable, Iterator
-
-from typing_extensions import ParamSpec, TypeAliasType, TypeVar
+from typing import Generic, Optional, cast
+
+from typing_extensions import (
+    ParamSpec,
+    Protocol,
+    TypeAliasType,
+    TypeVar,
+    override,
+    runtime_checkable,
+)
 
 from ..util import ContextResource, resource_factory_contextmanager
 from .utils import LovelyStats, format_tensor_stats
@@ -17,16 +25,26 @@ P = ParamSpec("P")
 
 LovelyStatsFn = TypeAliasType(
     "LovelyStatsFn",
-    Callable[[TArray], LovelyStats],
-    type_params=(TArray,),
-)
-LovelyReprFn = TypeAliasType(
-    "LovelyReprFn",
-    Callable[[TArray], str],
+    Callable[[TArray], Optional[LovelyStats]],
     type_params=(TArray,),
 )
 
 
+@runtime_checkable
+class LovelyReprFn(Protocol[TArray]):
+    @property
+    def __lovely_repr_instance__(self) -> lovely_repr[TArray]: ...
+
+    @__lovely_repr_instance__.setter
+    def __lovely_repr_instance__(self, value: lovely_repr[TArray]) -> None: ...
+
+    @property
+    def __name__(self) -> str: ...
+
+    def set_fallback_repr(self, repr_fn: Callable[[TArray], str]) -> None: ...
+    def __call__(self, value: TArray, /) -> str: ...
+
+
 def _find_missing_deps(dependencies: list[str]):
     missing_deps: list[str] = []
 
@@ -39,50 +57,61 @@ def _find_missing_deps(dependencies: list[str]):
     return missing_deps
 
 
-def lovely_repr(dependencies: list[str]):
-    """
-    Decorator to create a lovely representation function for an array.
-
-    Args:
-        dependencies: List of dependencies to check before running the function.
-            If any dependency is not available, the function will not run.
-
-    Returns:
-        A decorator function that takes a function and returns a lovely representation function.
-
-    Example:
-        @lovely_repr(dependencies=["torch"])
-        def my_array_stats(array):
-            return {...}
-    """
-
-    def decorator_fn(array_stats_fn: LovelyStatsFn[TArray]) -> LovelyReprFn[TArray]:
+class lovely_repr(Generic[TArray]):
+    @override
+    def __init__(
+        self,
+        dependencies: list[str],
+        fallback_repr: Callable[[TArray], str] | None = None,
+    ):
         """
         Decorator to create a lovely representation function for an array.
 
         Args:
-            array_stats_fn: A function that takes an array and returns its stats.
-
+            dependencies: List of dependencies to check before running the function.
+                If any dependency is not available, the function will not run.
+            fallback_repr: A function that takes an array and returns its fallback representation.
         Returns:
-            A function that takes an array and returns its lovely representation.
+            A decorator function that takes a function and returns a lovely representation function.
+
+        Example:
+            @lovely_repr(dependencies=["torch"])
+            def my_array_stats(array):
+                return {...}
         """
+        super().__init__()
+
+        if fallback_repr is None:
+            fallback_repr = repr
+
+        self._dependencies = dependencies
+        self._fallback_repr = fallback_repr
+
+    def set_fallback_repr(self, repr_fn: Callable[[TArray], str]) -> None:
+        self._fallback_repr = repr_fn
 
+    def __call__(
+        self, array_stats_fn: LovelyStatsFn[TArray], /
+    ) -> LovelyReprFn[TArray]:
         @functools.wraps(array_stats_fn)
-        def wrapper(array: TArray) -> str:
-            if missing_deps := _find_missing_deps(dependencies):
+        def wrapper_fn(array: TArray) -> str:
+            if missing_deps := _find_missing_deps(self._dependencies):
                 log.warning(
                     f"Missing dependencies: {', '.join(missing_deps)}. "
                     "Skipping lovely representation."
                 )
-                return repr(array)
+                return self._fallback_repr(array)
+
+            if (stats := array_stats_fn(array)) is None:
+                return self._fallback_repr(array)
 
-            stats = array_stats_fn(array)
             return format_tensor_stats(stats)
 
+        wrapper = cast(LovelyReprFn[TArray], wrapper_fn)
+        wrapper.__lovely_repr_instance__ = self
+        wrapper.set_fallback_repr = self.set_fallback_repr
         return wrapper
 
-    return decorator_fn
-
 
 LovelyMonkeyPatchInputFn = TypeAliasType(
     "LovelyMonkeyPatchInputFn",

{nshutils-0.22.5 → nshutils-0.30.1}/src/nshutils/lovely/jax_.py

@@ -54,9 +54,13 @@ def _device(array: jax.Array) -> str:
 
 
 @lovely_repr(dependencies=["jax"])
-def jax_repr(array: jax.Array) -> LovelyStats:
+def jax_repr(array: jax.Array) -> LovelyStats | None:
     import jax.numpy as jnp
 
+    # For dtypes like `object` or `str`, we let the fallback repr handle it
+    if not jnp.issubdtype(array.dtype, jnp.number):
+        return None
+
     return {
         # Basic attributes
         "shape": array.shape,
@@ -79,6 +83,7 @@ def jax_monkey_patch():
 
     prev_repr = array.ArrayImpl.__repr__
     prev_str = array.ArrayImpl.__str__
+    jax_repr.set_fallback_repr(prev_repr)
     try:
         patch_to(array.ArrayImpl, "__repr__", jax_repr)
         patch_to(array.ArrayImpl, "__str__", jax_repr)

{nshutils-0.22.5 → nshutils-0.30.1}/src/nshutils/lovely/numpy_.py

@@ -52,7 +52,11 @@ def _dtype_str(array: np.ndarray) -> str:
 
 
 @lovely_repr(dependencies=["numpy"])
-def numpy_repr(array: np.ndarray) -> LovelyStats:
+def numpy_repr(array: np.ndarray) -> LovelyStats | None:
+    # For dtypes like `object` or `str`, we let the fallback repr handle it
+    if not np.issubdtype(array.dtype, np.number):
+        return None
+
     return {
         # Basic attributes
         "shape": array.shape,
@@ -67,6 +71,9 @@ def numpy_repr(array: np.ndarray) -> LovelyStats:
     }
 
 
+numpy_repr.set_fallback_repr(np.array_repr)
+
+
 # If numpy 2.0, use the new API override_repr.
 if _np_ge_2():
 

{nshutils-0.22.5 → nshutils-0.30.1}/src/nshutils/lovely/torch_.py

@@ -60,7 +60,7 @@ def _to_np(tensor: torch.Tensor) -> np.ndarray:
 
 
 @lovely_repr(dependencies=["torch"])
-def torch_repr(tensor: torch.Tensor) -> LovelyStats:
+def torch_repr(tensor: torch.Tensor) -> LovelyStats | None:
     return {
         # Basic attributes
         "shape": tensor.shape,
@@ -87,10 +87,15 @@ def torch_monkey_patch():
     original_repr = torch.Tensor.__repr__
     original_str = torch.Tensor.__str__
     original_parameter_repr = torch.nn.Parameter.__repr__
+    torch_repr.set_fallback_repr(original_repr)
+
     try:
         patch_to(torch.Tensor, "__repr__", torch_repr)
         patch_to(torch.Tensor, "__str__", torch_repr)
-        del torch.nn.Parameter.__repr__
+        try:
+            delattr(torch.nn.Parameter, "__repr__")
+        except AttributeError:
+            pass
 
         yield
     finally: