nshutils 0.31.1__py3-none-any.whl → 0.32.0__py3-none-any.whl

nshutils/actsave/_saver.py

@@ -20,7 +20,7 @@ from ..collections import apply_to_collection
 
 if not TYPE_CHECKING:
     try:
-        import torch  # type: ignore
+        import torch  # pyright: ignore[reportMissingImports]
 
         Tensor = torch.Tensor
         _torch_installed = True
@@ -30,9 +30,9 @@ if not TYPE_CHECKING:
 
         Tensor = Never
 else:
-    import torch  # type: ignore
+    import torch  # pyright: ignore[reportMissingImports]
 
-    Tensor = torch.Tensor
+    Tensor = TypeAliasType("Tensor", torch.Tensor)
     _torch_installed: Literal[True] = True
 
 log = getLogger(__name__)
@@ -60,7 +60,7 @@ def _to_numpy(activation: Value) -> np.ndarray:
         return np.array(activation)
     elif isinstance(activation, np.ndarray):
         return activation
-    elif _torch_installed and isinstance(activation, Tensor):
+    elif _torch_installed and isinstance(activation, torch.Tensor):
         activation_ = activation.detach()
         if activation_.is_floating_point():
             # NOTE: We need to convert to float32 because [b]float16 is not supported by numpy
@@ -126,7 +126,8 @@ class Activation:
         if activation is None:
             return None
 
-        activation = apply_to_collection(activation, Tensor, _to_numpy)
+        if _torch_installed:
+            activation = apply_to_collection(activation, Tensor, _to_numpy)
         activation = _to_numpy(activation)
 
         # Set the transformed value

nshutils/lovely/_monkey_patch_all.py

@@ -30,9 +30,10 @@ def _find_deps() -> list[Library]:
 
 class monkey_patch(lovely_patch):
     def __init__(self, libraries: list[Library] | Literal["auto"] = "auto"):
-        self.libraries = libraries
-        if self.libraries == "auto":
+        if libraries == "auto":
             self.libraries = _find_deps()
+        else:
+            self.libraries = libraries
 
         if not self.libraries:
             raise ValueError(
@@ -59,7 +60,7 @@ class monkey_patch(lovely_patch):
 
                 self.stack.enter_context(numpy_monkey_patch())
             else:
-                assert_never(library)  # type: ignore
+                assert_never(library)
 
         log.info(
             f"Monkey patched libraries: {', '.join(self.libraries)}. "

nshutils/lovely/jax_.py

@@ -9,7 +9,7 @@ from ._base import lovely_patch, lovely_repr
 from .utils import LovelyStats, array_stats, patch_to
 
 if TYPE_CHECKING:
-    import jax
+    import jax  # pyright: ignore[reportMissingImports]
 
 
 def _type_name(array: jax.Array):
@@ -42,7 +42,7 @@ def _dtype_str(array: jax.Array) -> str:
 
 
 def _device(array: jax.Array) -> str:
-    from jaxlib.xla_extension import Device
+    from jaxlib.xla_extension import Device  # pyright: ignore[reportMissingImports]
 
     if callable(device := array.device):
         device = device()
@@ -56,7 +56,7 @@ def _device(array: jax.Array) -> str:
 
 @lovely_repr(dependencies=["jax"])
 def jax_repr(array: jax.Array) -> LovelyStats | None:
-    import jax.numpy as jnp
+    import jax.numpy as jnp  # pyright: ignore[reportMissingImports]
 
     # For dtypes like `object` or `str`, we let the fallback repr handle it
     if not jnp.issubdtype(array.dtype, jnp.number):
@@ -85,7 +85,7 @@ class jax_monkey_patch(lovely_patch):
 
     @override
     def patch(self):
-        from jax._src import array
+        from jax._src import array  # pyright: ignore[reportMissingImports]
 
         self.prev_repr = array.ArrayImpl.__repr__
         self.prev_str = array.ArrayImpl.__str__
@@ -96,7 +96,7 @@ class jax_monkey_patch(lovely_patch):
 
     @override
     def unpatch(self):
-        from jax._src import array
+        from jax._src import array  # pyright: ignore[reportMissingImports]
 
         patch_to(array.ArrayImpl, "__repr__", self.prev_repr)
         patch_to(array.ArrayImpl, "__str__", self.prev_str)

nshutils/lovely/torch_.py

@@ -9,11 +9,11 @@ from ._base import lovely_patch, lovely_repr
 from .utils import LovelyStats, array_stats, patch_to
 
 if TYPE_CHECKING:
-    import torch
+    import torch  # pyright: ignore[reportMissingImports]
 
 
 def _type_name(tensor: torch.Tensor):
-    import torch
+    import torch  # pyright: ignore[reportMissingImports]
 
     return (
         "tensor"
@@ -45,7 +45,7 @@ def _dtype_str(tensor: torch.Tensor) -> str:
 
 
 def _to_np(tensor: torch.Tensor) -> np.ndarray:
-    import torch
+    import torch  # pyright: ignore[reportMissingImports]
 
     # Get tensor data as CPU NumPy array for analysis
     t_cpu = tensor.detach().cpu()
@@ -88,7 +88,7 @@ class torch_monkey_patch(lovely_patch):
 
     @override
     def patch(self):
-        import torch
+        import torch  # pyright: ignore[reportMissingImports]
 
         self.original_repr = torch.Tensor.__repr__
         self.original_str = torch.Tensor.__str__
@@ -104,7 +104,7 @@ class torch_monkey_patch(lovely_patch):
 
     @override
     def unpatch(self):
-        import torch
+        import torch  # pyright: ignore[reportMissingImports]
 
         patch_to(torch.Tensor, "__repr__", self.original_repr)
         patch_to(torch.Tensor, "__str__", self.original_str)

nshutils/snoop.py

@@ -20,19 +20,24 @@ try:
     import warnings
     from contextlib import nullcontext
 
-    import pysnooper  # type: ignore
-    import pysnooper.utils  # type: ignore
+    import pysnooper  # pyright: ignore[reportMissingImports]
+    import pysnooper.utils  # pyright: ignore[reportMissingImports]
 
     try:
-        import torch  # type: ignore
+        import torch  # pyright: ignore[reportMissingImports]
     except ImportError:
         torch = None
 
     try:
-        import numpy  # type: ignore
+        import numpy  # pyright: ignore[reportMissingImports]
     except ImportError:
         numpy = None
 
+    try:
+        import jax  # pyright: ignore[reportMissingImports]
+    except ImportError:
+        jax = None
+
     FLOATING_POINTS = set()
     for i in ["float", "double", "half", "complex128", "complex32", "complex64"]:
         # older version of PyTorch do not have complex dtypes
@@ -48,17 +53,25 @@ try:
 
     def default_format(x):
         try:
-            import lovely_tensors as lt  # type: ignore
+            from .lovely import torch_repr
 
-            return str(lt.lovely(x))
+            return torch_repr(x)
         except BaseException:
             return str(x.shape)
 
     def default_numpy_format(x):
         try:
-            import lovely_numpy as lo  # type: ignore
+            from .lovely import numpy_repr
+
+            return numpy_repr(x)
+        except BaseException:
+            return str(x.shape)
+
+    def default_jax_format(x):
+        try:
+            from .lovely import jax_repr
 
-            return str(lo.lovely(x))
+            return jax_repr(x)
         except BaseException:
             return str(x.shape)
 
@@ -68,6 +81,7 @@ try:
             *args,
             tensor_format=default_format,
             numpy_format=default_numpy_format,
+            jax_format=default_jax_format,
             **kwargs,
         ):
             self.orig_custom_repr = (
@@ -78,6 +92,7 @@ try:
             super(TorchSnooper, self).__init__(*args, **kwargs)
             self.tensor_format = tensor_format
             self.numpy_format = numpy_format
+            self.jax_format = jax_format
 
         @staticmethod
         def is_return_types(x):
@@ -176,6 +191,8 @@ try:
                 return self.tensor_format(x)
             if numpy is not None and isinstance(x, numpy.ndarray):
                 return self.numpy_format(x)
+            if jax is not None and isinstance(x, jax.Array):
+                return self.jax_format(x)
             if self.is_return_types(x):
                 return self.return_types_repr(x)
             if orig_repr_func is not repr:

nshutils/typecheck.py

@@ -38,18 +38,18 @@ from jaxtyping._storage import get_shape_memo, shape_str
 from typing_extensions import TypeVar
 
 try:
-    import torch  # type: ignore
+    import torch  # pyright: ignore[reportMissingImports]
 except ImportError:
     torch = None
 
 try:
-    import np  # type: ignore
+    import np  # pyright: ignore[reportMissingImports]
 except ImportError:
     np = None
 
 
 try:
-    import jax  # type: ignore
+    import jax  # pyright: ignore[reportMissingImports]
 except ImportError:
     jax = None
 
@@ -124,23 +124,23 @@ def _make_error_str(input: Any, t: Any) -> str:
         error_components.append(t.__instancecheck_str__(input))
     if torch is not None and torch.is_tensor(input):
         try:
-            from lovely_tensors import lovely  # type: ignore
+            from .lovely import torch_repr
 
-            error_components.append(repr(lovely(input)))
+            error_components.append(torch_repr(input))
         except BaseException:
             error_components.append(repr(input.shape))
     elif jax is not None and isinstance(input, jax.Array):
         try:
-            from lovely_jax import lovely  # type: ignore
+            from .lovely import jax_repr
 
-            error_components.append(repr(lovely(input)))
+            error_components.append(jax_repr(input))
         except BaseException:
             error_components.append(repr(input.shape))
     elif np is not None and isinstance(input, np.ndarray):
         try:
-            from lovely_numpy import lovely  # type: ignore
+            from .lovely import numpy_repr
 
-            error_components.append(repr(lovely(input)))
+            error_components.append(numpy_repr(input))
         except BaseException:
             error_components.append(repr(input.shape))
     error_components.append(shape_str(get_shape_memo()))

{nshutils-0.31.1.dist-info → nshutils-0.32.0.dist-info}/METADATA

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

{nshutils-0.31.1.dist-info → nshutils-0.32.0.dist-info}/RECORD

@@ -2,20 +2,20 @@ nshutils/__init__.py,sha256=AFx1d5k34MyJ2kCHQL5vrZB8GDp2nYUaIUEjszSa25I,477
 nshutils/__init__.pyi,sha256=R4TIk--jAgVyTibdgezJQTMce3HpMCNakAJeaDqA6bc,676
 nshutils/actsave/__init__.py,sha256=hAVsog9d1g3_rQN1TRslrl6sK1PhCGbjy8PPUAmJI58,203
 nshutils/actsave/_loader.py,sha256=btLSQdErpTmK6VyG8PxJrJNsztzyavSF71n4Ec3_49E,7619
-nshutils/actsave/_saver.py,sha256=_qkX0NZYvy31hdlyfhneac4kUNS_44XjOG0ZtKpdqrg,12720
+nshutils/actsave/_saver.py,sha256=zcsuAP7JhaXy3EzBZillyiIf1EeIsoilNoDXWKvDta0,12834
 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=684eZOKLmSgsTcCVlWk1Ip1cxJxmz-rKeXLmWXuCEWA,487
 nshutils/lovely/_base.py,sha256=kJY-UhdFTRBlAg_YzfJmG4ICb6vSdOJKiRc6vksxvoE,4424
-nshutils/lovely/_monkey_patch_all.py,sha256=xq09InGcOsGDrELV_KIrhE0H4EWyMdrUZ_1_BR2e_b0,2224
+nshutils/lovely/_monkey_patch_all.py,sha256=dgfqrJU1sYY7icV_yupG1gMdgTHo-Mn9EeiWVxcfq9I,2221
 nshutils/lovely/config.py,sha256=lVNMuU1oUvsYlGN0Sn-m6iOLbJIchVnWDpyHm09nWo8,1224
-nshutils/lovely/jax_.py,sha256=PGnv33LrEM3aLvXLBbAx4b7dOkJwONidyPZjToZ62Og,2592
+nshutils/lovely/jax_.py,sha256=Hju21e2vjc2ps6tu6yXiR_GAq03S16WLUzcITrHtk60,2797
 nshutils/lovely/numpy_.py,sha256=BBP9663l4Hr-TB34xDMHQQZ1zpuOgBegUOGl7_wV6R0,3503
-nshutils/lovely/torch_.py,sha256=J1pDJY1zzEANqa6EaJpG1pc_SYgM8YWOo1TjWdVeiA0,2946
+nshutils/lovely/torch_.py,sha256=4r46p4tTaT6cWHl1_VcHh6x3RBST4xijqF2h7f9fnVg,3151
 nshutils/lovely/utils.py,sha256=2ksT5YGVViFuWc8jSkwVCsABripJmyVJdEDDH7aab70,10459
-nshutils/snoop.py,sha256=7d7_Q5sJmINL1J29wcnxEvpV95zvZYNoVn5frCq-rww,7393
-nshutils/typecheck.py,sha256=Gi7xtfilN_UwZ1FTFqBVKDhcQzBEDonVxIv3bUj-uXY,5582
-nshutils-0.31.1.dist-info/METADATA,sha256=pW-XE-rF3TPtok9VPan8cG6IlaEjlX55NBL4SKaNOHQ,4406
-nshutils-0.31.1.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-nshutils-0.31.1.dist-info/RECORD,,
+nshutils/snoop.py,sha256=ZkqodwLOey3cAIZBK_OIV4xT7PyU10Y4L2oApi_87k8,7941
+nshutils/typecheck.py,sha256=U6QDrAwxGtoQW5GYcifWJQkUicCz2iutPw1kexGPKaU,5596
+nshutils-0.32.0.dist-info/METADATA,sha256=Hl3SklKnO3BbyZORJTdMxXTaBUmdnYB3KW0r02ucHM8,4406
+nshutils-0.32.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+nshutils-0.32.0.dist-info/RECORD,,