nshutils 0.10.0__py3-none-any.whl → 0.11.0__py3-none-any.whl

nshutils/actsave/_saver.py

@@ -1,7 +1,6 @@
 import contextlib
 import fnmatch
 import tempfile
-import uuid
 import weakref
 from collections.abc import Callable, Mapping
 from dataclasses import dataclass
@@ -12,6 +11,7 @@ from typing import TYPE_CHECKING, Generic, TypeAlias, cast, overload
 
 import numpy as np
 from typing_extensions import Never, ParamSpec, TypeVar, override
+from uuid_extensions import uuid7str
 
 from ..collections import apply_to_collection
 
@@ -236,21 +236,32 @@ class ActSaveProvider:
     _saver: _Saver | None = None
     _prefixes: list[str] = []
 
-    def initialize(self, save_dir: Path | None = None):
+    def enable(self, save_dir: Path | None = None):
         """
         Initializes the saver with the given configuration and save directory.
 
         Args:
             save_dir (Path): The directory where the saved files will be stored.
         """
+        if self._saver is not None:
+            log.warning("ActSave is already enabled")
+            return
+
+        if save_dir is None:
+            save_dir = Path(tempfile.gettempdir()) / f"actsave-{uuid7str()}"
+            log.critical(f"No save_dir specified, using {save_dir=}")
+        self._saver = _Saver(save_dir, lambda: self._prefixes)
+
+    def disable(self):
+        """
+        Disables the actsaver.
+        """
         if self._saver is None:
-            if save_dir is None:
-                save_dir = Path(tempfile.gettempdir()) / f"actsave-{uuid.uuid4()}"
-                log.critical(f"No save_dir specified, using {save_dir=}")
-            self._saver = _Saver(
-                save_dir,
-                lambda: self._prefixes,
-            )
+            log.warning("ActSave is already disabled")
+            return
+
+        del self._saver
+        self._saver = None
 
     @contextlib.contextmanager
     def enabled(self, save_dir: Path | None = None):
@@ -260,12 +271,16 @@ class ActSaveProvider:
         Args:
             save_dir (Path): The directory where the saved files will be stored.
         """
-        prev = self._saver
-        self.initialize(save_dir)
+        if self._saver is not None:
+            log.warning("ActSave is already enabled")
+            yield
+            return
+
+        self.enable(save_dir)
         try:
             yield
         finally:
-            self._saver = prev
+            self.disable()
 
     @override
     def __init__(self):

nshutils-0.11.0.dist-info/METADATA

@@ -0,0 +1,120 @@
+Metadata-Version: 2.1
+Name: nshutils
+Version: 0.11.0
+Summary: 
+Author: Nima Shoghi
+Author-email: nimashoghi@gmail.com
+Requires-Python: >=3.10,<4.0
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Provides-Extra: extra
+Requires-Dist: beartype
+Requires-Dist: jaxtyping
+Requires-Dist: numpy
+Requires-Dist: pysnooper ; extra == "extra"
+Requires-Dist: typing-extensions
+Requires-Dist: uuid7
+Description-Content-Type: text/markdown
+
+# nshutils
+
+`nshutils` is a collection of utility functions and classes that I've found useful in my day-to-day work as an ML researcher. This library includes utilities for typechecking, logging, and saving/loading activations from neural networks.
+
+## Installation
+
+To install `nshutils`, simply run:
+
+```bash
+pip install nshutils
+```
+
+## Features
+
+### Typechecking
+
+`nshutils` provides a simple way to typecheck your code using the `jaxtyping` library. Simply call `typecheck_this_module()` at the top of your module (i.e., in the root `__init__.py` file) to enable typechecking for the entire module:
+
+```python
+from nshutils.typecheck import typecheck_this_module
+
+typecheck_this_module()
+```
+
+You can also use the `tassert` function to assert that a value is of a certain type:
+
+```python
+from nshutils.typecheck import tassert, Float32
+
+def my_function(x: Float32):
+    tassert(Float32, x)
+    ...
+```
+
+### Logging
+
+`nshutils` provides a simple way to configure logging for your project. Simply call one of the logging setup functions:
+
+```python
+from nshutils.logging import init_python_logging
+
+init_python_logging()
+```
+
+This will configure logging to use pretty formatting for PyTorch tensors and numpy arrays (inspired by and/or utilizing [`lovely-numpy`](https://github.com/xl0/lovely-numpy) and [`lovely-tensors`](https://github.com/xl0/lovely-tensors)), and will also enable rich logging if the `rich` library is installed.
+
+### Activation Saving/Loading
+
+`nshutils` provides a simple way to save and load activations from neural networks. To save activations, use the `ActSave` object:
+
+```python
+from nshutils import ActSave
+
+def my_model_forward(x):
+    ...
+    # Save activations to "{save_dir}/encoder.activations/{idx}.npy"
+    ActSave({"encoder.activations": x})
+
+    # Equivalent to the above
+    with ActSave.context("encoder"):
+        ActSave(activations=x)
+    ...
+
+ActSave.enable(save_dir="path/to/activations")
+x = torch.randn(...)
+my_model_forward(x)
+# Activations are saved to disk under the "path/to/activations" directory
+```
+
+This will save the `x` tensor to disk under the `encoder` prefix.
+
+To load activations, use the `ActLoad` class:
+
+```python
+from nshutils import ActLoad
+
+act_load = ActLoad.from_latest_version("path/to/activations")
+encoder_acts = act_load["encoder"]
+
+for act in encoder_acts:
+    print(act.shape)
+```
+
+This will load all of the activations saved under the `encoder` prefix.
+
+### Other Utilities
+
+`nshutils` also provides a few other utility functions/classes:
+
+- `snoop`: A simple way to debug your code using the `pysnooper` library, based on the [`torchsnooper`](https://github.com/zasdfgbnm/TorchSnooper) library.
+- `apply_to_collection`: Recursively apply a function to all elements of a collection that match a certain type, taken from the [`pytorch-lightning`](https://github.com/Lightning-AI/pytorch-lightning) library.
+
+## Contributing
+
+Contributions to `nshutils` are welcome! Please open an issue or submit a pull request on the [GitHub repository](https://github.com/nimashoghi/nshutils).
+
+## License
+
+`nshutils` is released under the MIT License. See the `LICENSE` file for more details.
+

{nshutils-0.10.0.dist-info → nshutils-0.11.0.dist-info}/RECORD

@@ -1,11 +1,11 @@
 nshutils/__init__.py,sha256=ZRdp3KcreswdrusXOGc8aEI07qfBCYQGJYC7NqZ5WPI,324
 nshutils/actsave/__init__.py,sha256=6gJ49011Ad3kS8BejeZRPj7ZyVyIcmX-VKLEmYZyGM8,167
 nshutils/actsave/_loader.py,sha256=fAhD32DrJa4onkYfcwc21YIeGEYzOSXCK_HVo9SZLgQ,4604
-nshutils/actsave/_saver.py,sha256=M7PwFy6faefEDG9hbclbWUNGi1yX7ZjzshdFdmsrtyk,9890
+nshutils/actsave/_saver.py,sha256=-uKMmMKjEErCa8pfLhAW4077GFcUjxr0Qq0PFAnbEOw,10254
 nshutils/collections.py,sha256=EE_qLd-LrsX5lsyk9GSKh03Q8bhn9CHB3jiEeNj4uF4,5197
 nshutils/logging.py,sha256=tL-6XvdvJEEr7bje9DSmUBpm6pnJS9XG_1fuw3U1eME,2573
 nshutils/snoop.py,sha256=2RNlOcBFnCUyOsfGlbBb8FwSGq03hxfypMfqMA0No-A,7356
 nshutils/typecheck.py,sha256=_KtfinRy9A0Dgq78kN5MGGrefvb6jn2tGY6svdLEBAs,4915
-nshutils-0.10.0.dist-info/METADATA,sha256=QlUQG8pRWMUlSJP56NKqRqGfkc6uv8VOsMj8RzE-Kq8,559
-nshutils-0.10.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshutils-0.10.0.dist-info/RECORD,,
+nshutils-0.11.0.dist-info/METADATA,sha256=KIoZDV9c0bjTKbV7F-YOxMthSo_a9J5KOStw92OTa7s,3741
+nshutils-0.11.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshutils-0.11.0.dist-info/RECORD,,

nshutils-0.10.0.dist-info/METADATA

@@ -1,20 +0,0 @@
-Metadata-Version: 2.1
-Name: nshutils
-Version: 0.10.0
-Summary: 
-Author: Nima Shoghi
-Author-email: nimashoghi@gmail.com
-Requires-Python: >=3.10,<4.0
-Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.10
-Classifier: Programming Language :: Python :: 3.11
-Classifier: Programming Language :: Python :: 3.12
-Provides-Extra: extra
-Requires-Dist: beartype
-Requires-Dist: jaxtyping
-Requires-Dist: numpy
-Requires-Dist: pysnooper ; extra == "extra"
-Requires-Dist: typing-extensions
-Description-Content-Type: text/markdown
-
-