nshutils 0.22.6__py3-none-any.whl → 0.30.1__py3-none-any.whl

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/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 | None],
-    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,53 +57,61 @@ def _find_missing_deps(dependencies: list[str]):
     return missing_deps
 
 
-def lovely_repr(dependencies: list[str], fallback_repr: Callable[[TArray], 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.
-        fallback_repr: A function that takes an array and returns its fallback representation.
-    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,
-                or `None` if the array is not supported.
-
+            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 fallback_repr(array)
+                return self._fallback_repr(array)
 
             if (stats := array_stats_fn(array)) is None:
-                return fallback_repr(array)
+                return self._fallback_repr(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/lovely/jax_.py

@@ -53,7 +53,7 @@ def _device(array: jax.Array) -> str:
     return f"{device.platform}:{device.id}"
 
 
-@lovely_repr(dependencies=["jax"], fallback_repr=jax.Array.__repr__)
+@lovely_repr(dependencies=["jax"])
 def jax_repr(array: jax.Array) -> LovelyStats | None:
     import jax.numpy as jnp
 
@@ -83,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/lovely/numpy_.py

@@ -51,7 +51,7 @@ def _dtype_str(array: np.ndarray) -> str:
     return dtype_base
 
 
-@lovely_repr(dependencies=["numpy"], fallback_repr=np.array_repr)
+@lovely_repr(dependencies=["numpy"])
 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):
@@ -71,6 +71,9 @@ def numpy_repr(array: np.ndarray) -> LovelyStats | None:
     }
 
 
+numpy_repr.set_fallback_repr(np.array_repr)
+
+
 # If numpy 2.0, use the new API override_repr.
 if _np_ge_2():
 

nshutils/lovely/torch_.py

@@ -59,7 +59,7 @@ def _to_np(tensor: torch.Tensor) -> np.ndarray:
     return t_np
 
 
-@lovely_repr(dependencies=["torch"], fallback_repr=torch.Tensor.__repr__)
+@lovely_repr(dependencies=["torch"])
 def torch_repr(tensor: torch.Tensor) -> LovelyStats | None:
     return {
         # Basic attributes
@@ -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:

{nshutils-0.22.6.dist-info → nshutils-0.30.1.dist-info}/METADATA

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

{nshutils-0.22.6.dist-info → nshutils-0.30.1.dist-info}/RECORD

@@ -2,21 +2,21 @@ nshutils/__init__.py,sha256=AFx1d5k34MyJ2kCHQL5vrZB8GDp2nYUaIUEjszSa25I,477
 nshutils/__init__.pyi,sha256=ICbY2_XBAlXIVOGyK4PQpatmlUFHHc5-bqM4sfFZoAY,613
 nshutils/actsave/__init__.py,sha256=hAVsog9d1g3_rQN1TRslrl6sK1PhCGbjy8PPUAmJI58,203
 nshutils/actsave/_loader.py,sha256=mof3HezUNvLliz7macstX6ewXW05L0Mtv3zJyrbmImg,4640
-nshutils/actsave/_saver.py,sha256=GyoJpeJIIG_Y2Hr3MTXylTM6YNyoInf2mvKs0foU7no,10528
+nshutils/actsave/_saver.py,sha256=IS9TVP8WUizoj5fHrQ6hodtjidT__LDRwz5aoWHupVo,12013
 nshutils/collections.py,sha256=QWGyANmo4Efq4XRNHDSTE9tRLStwEZHGwE0ATHR-Vqo,5233
 nshutils/display.py,sha256=Ge63yllx7gi-MKL3mKQeQ5doql_nj56-o5aoTVmusDg,1473
 nshutils/logging.py,sha256=78pv3-I_gmbKSf5_mYYBr6_H4GNBGErghAdhH9wfYIc,2205
 nshutils/lovely/__init__.py,sha256=gbWMNs7xfK1CiNdkHvfH0KcyaGjdZ8_WUBGfaEUDN4I,451
-nshutils/lovely/_base.py,sha256=SH8MY6wbucB7us4pU1cekoU2glanutEB8vDrLDkLEL0,4544
+nshutils/lovely/_base.py,sha256=-JYF2zci04PJjmkBdm_iV3uWgD_d7e5zCIAINDlQIKc,5266
 nshutils/lovely/_monkey_patch_all.py,sha256=zgMupp2Wc_O9R3arl-BAIePpvQSi6TCeshGMaui-Cc8,1986
 nshutils/lovely/config.py,sha256=lVNMuU1oUvsYlGN0Sn-m6iOLbJIchVnWDpyHm09nWo8,1224
-nshutils/lovely/jax_.py,sha256=J64CtP2mx_yPQ7ZiLKSrqEJfPDBBbx3j6o4zKc1f1eg,2404
-nshutils/lovely/numpy_.py,sha256=iaTUQpamudyWd7pTBJyj_hcS9X8H53muCw2IYY9ZcUM,3345
-nshutils/lovely/torch_.py,sha256=Jt0c7ysfrbgxORhL6fpbxCnfBDoYaTtTyc935yCJ7H0,2682
+nshutils/lovely/jax_.py,sha256=c_hvlch_c9OZ0WJjFIeY46kKQcCELspwdmoexkKLsCg,2412
+nshutils/lovely/numpy_.py,sha256=GDOOuhCYfShfKUZiuI8J91eAm27urrYyxETTR-Mxz0E,3362
+nshutils/lovely/torch_.py,sha256=9diSkM1L2B6l0yQqTRBoZUVElyqgHUcJdFCXD3NvTxk,2767
 nshutils/lovely/utils.py,sha256=2ksT5YGVViFuWc8jSkwVCsABripJmyVJdEDDH7aab70,10459
 nshutils/snoop.py,sha256=7d7_Q5sJmINL1J29wcnxEvpV95zvZYNoVn5frCq-rww,7393
 nshutils/typecheck.py,sha256=Gi7xtfilN_UwZ1FTFqBVKDhcQzBEDonVxIv3bUj-uXY,5582
 nshutils/util.py,sha256=tx-XiRbOrpafV3OkJDE5IVFtzn3kN7uSZ8FkMor0H5c,2845
-nshutils-0.22.6.dist-info/METADATA,sha256=cclgVNUapsycUT65AJHSQXQpej_WEwwhRIl9vwrLMbY,4406
-nshutils-0.22.6.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-nshutils-0.22.6.dist-info/RECORD,,
+nshutils-0.30.1.dist-info/METADATA,sha256=iFnO4L_bpOdtCm8f-m7IvA4A8WtkWb5T07dYi-aUnzI,4406
+nshutils-0.30.1.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+nshutils-0.30.1.dist-info/RECORD,,