nshutils 0.7.0__tar.gz → 0.10.0__tar.gz

{nshutils-0.7.0 → nshutils-0.10.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshutils
-Version: 0.7.0
+Version: 0.10.0
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com
@@ -9,10 +9,11 @@ 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
+Requires-Dist: pysnooper ; extra == "extra"
 Requires-Dist: typing-extensions
 Description-Content-Type: text/markdown
 

{nshutils-0.7.0 → nshutils-0.10.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nshutils"
-version = "0.7.0"
+version = "0.10.0"
 description = ""
 authors = ["Nima Shoghi <nimashoghi@gmail.com>"]
 readme = "README.md"
@@ -8,10 +8,10 @@ readme = "README.md"
 [tool.poetry.dependencies]
 python = "^3.10"
 numpy = "*"
-pysnooper = "*"
 jaxtyping = "*"
 typing-extensions = "*"
 beartype = "*"
+pysnooper = { version = "*", optional = true }
 
 
 [tool.poetry.group.dev.dependencies]
@@ -31,7 +31,9 @@ strictListInference = true
 strictDictionaryInference = true
 strictSetInference = true
 reportPrivateImportUsage = false
-ignore = ["./build/"]
 
 [tool.ruff.lint]
 ignore = ["F722", "F821", "E731", "E741"]
+
+[tool.poetry.extras]
+extra = ["pysnooper"]

{nshutils-0.7.0 → nshutils-0.10.0}/src/nshutils/__init__.py

@@ -1,5 +1,7 @@
 from . import actsave as actsave
 from . import typecheck as typecheck
+from .actsave import ActLoad as ActLoad
+from .actsave import ActSave as ActSave
 from .logging import init_python_logging as init_python_logging
 from .logging import lovely as lovely
 from .logging import pretty as pretty

{nshutils-0.7.0 → nshutils-0.10.0}/src/nshutils/actsave/__init__.py

@@ -1,6 +1,4 @@
-from ._loader import ActivationLoader as ActivationLoader
 from ._loader import ActLoad as ActLoad
 from ._saver import Activation as Activation
-from ._saver import ActivationSaver as ActivationSaver
 from ._saver import ActSave as ActSave
 from ._saver import Transform as Transform

{nshutils-0.7.0 → nshutils-0.10.0}/src/nshutils/actsave/_saver.py

@@ -53,11 +53,11 @@ def _to_numpy(activation: Value) -> np.ndarray:
     elif isinstance(activation, np.ndarray):
         return activation
     elif isinstance(activation, Tensor):
-        activation = activation.detach()
-        if activation.is_floating_point():
+        activation_ = activation.detach()
+        if activation_.is_floating_point():
             # NOTE: We need to convert to float32 because [b]float16 is not supported by numpy
-            activation = activation.float()
-        return activation.cpu().numpy()
+            activation_ = activation_.float()
+        return activation_.cpu().numpy()
     else:
         log.warning(f"Unrecognized activation type {type(activation)}")
 
@@ -358,4 +358,3 @@ class ActSaveProvider:
 
 
 ActSave = ActSaveProvider()
-ActivationSaver = ActSave