nshutils 0.6.1__py3-none-any.whl → 0.8.0__py3-none-any.whl

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/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/actsave/_saver.py

@@ -15,24 +15,25 @@ from typing_extensions import Never, ParamSpec, TypeVar, override
 
 from ..collections import apply_to_collection
 
-try:
-    import torch  # type: ignore
+if not TYPE_CHECKING:
+    try:
+        import torch  # type: ignore
 
-    if not TYPE_CHECKING:
         Tensor: TypeAlias = torch.Tensor
-except ImportError:
-    torch = None
+    except ImportError:
+        torch = None
 
-    if not TYPE_CHECKING:
         Tensor: TypeAlias = Never
+else:
+    import torch  # type: ignore
+
+    Tensor: TypeAlias = torch.Tensor
 
-if TYPE_CHECKING:
-    Tensor: TypeAlias = Never
 
 log = getLogger(__name__)
 
 Value: TypeAlias = int | float | complex | bool | str | np.ndarray | Tensor | None
-ValueOrLambda = Value | Callable[..., Value]
+ValueOrLambda: TypeAlias = Value | Callable[..., Value]
 
 
 def _torch_is_scripting() -> bool:
@@ -52,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)}")
 
@@ -357,4 +358,3 @@ class ActSaveProvider:
 
 
 ActSave = ActSaveProvider()
-ActivationSaver = ActSave

{nshutils-0.6.1.dist-info → nshutils-0.8.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshutils
-Version: 0.6.1
+Version: 0.8.0
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com
@@ -9,10 +9,10 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
-Requires-Dist: beartype (>=0.18.5,<0.19.0)
-Requires-Dist: jaxtyping (>=0.2.33,<0.3.0)
+Requires-Dist: beartype
+Requires-Dist: jaxtyping
 Requires-Dist: numpy
-Requires-Dist: pysnooper (>=1.2.0,<2.0.0)
+Requires-Dist: pysnooper
 Requires-Dist: typing-extensions
 Description-Content-Type: text/markdown
 

nshutils-0.8.0.dist-info/RECORD

@@ -0,0 +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/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.8.0.dist-info/METADATA,sha256=X5yqqnQA1Pin9-1VvhClw3l7YggpoMAZVUtg3l77LGk,517
+nshutils-0.8.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshutils-0.8.0.dist-info/RECORD,,

nshutils-0.6.1.dist-info/RECORD

@@ -1,11 +0,0 @@
-nshutils/__init__.py,sha256=1hxTdQ-e1ym76tdrSAaYoU4Qq5IUt1DRShi7khn7vs8,244
-nshutils/actsave/__init__.py,sha256=kafzChiViTeOY28GnwKD6yBJfiB7aSU8t1u5cZicdmc,280
-nshutils/actsave/_loader.py,sha256=fAhD32DrJa4onkYfcwc21YIeGEYzOSXCK_HVo9SZLgQ,4604
-nshutils/actsave/_saver.py,sha256=xpvnMi4dI8dFwu0tIpJVIp1S4J6FJT5fVYAUTWgvB7A,9885
-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.6.1.dist-info/METADATA,sha256=N3aKMAQ04xiBd9NvEYlV848JIG5zWjIZgZC20F7IQm0,571
-nshutils-0.6.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshutils-0.6.1.dist-info/RECORD,,