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

nshutils/actsave/_saver.py

@@ -9,7 +9,7 @@ from dataclasses import dataclass
 from functools import wraps
 from logging import getLogger
 from pathlib import Path
-from typing import TYPE_CHECKING, Generic, Literal, Union, cast, overload
+from typing import TYPE_CHECKING, Any, Generic, Literal, Union, cast, overload
 
 import numpy as np
 from typing_extensions import Never, ParamSpec, TypeAliasType, TypeVar, override
@@ -36,8 +36,9 @@ else:
 
 log = getLogger(__name__)
 
+# Updated to include Any for arbitrary types
 Value = TypeAliasType(
-    "Value", Union[int, float, complex, bool, str, np.ndarray, Tensor, None]
+    "Value", Union[int, float, complex, bool, str, np.ndarray, Tensor, Any, None]
 )
 ValueOrLambda = TypeAliasType("ValueOrLambda", Union[Value, Callable[..., Value]])
 
@@ -65,9 +66,8 @@ def _to_numpy(activation: Value) -> np.ndarray:
             activation_ = activation_.float()
         return activation_.cpu().numpy()
     else:
-        log.warning(f"Unrecognized activation type {type(activation)}")
-
-    return activation
+        # Handle arbitrary objects using numpy object dtype
+        return np.array(activation, dtype=object)
 
 
 T = TypeVar("T", infer_variance=True)

nshutils/lovely/_base.py

@@ -17,7 +17,7 @@ P = ParamSpec("P")
 
 LovelyStatsFn = TypeAliasType(
     "LovelyStatsFn",
-    Callable[[TArray], LovelyStats],
+    Callable[[TArray], LovelyStats | None],
     type_params=(TArray,),
 )
 LovelyReprFn = TypeAliasType(
@@ -39,14 +39,14 @@ def _find_missing_deps(dependencies: list[str]):
     return missing_deps
 
 
-def lovely_repr(dependencies: list[str]):
+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.
 
@@ -61,7 +61,8 @@ def lovely_repr(dependencies: list[str]):
         Decorator to create a lovely representation function for an array.
 
         Args:
-            array_stats_fn: A function that takes an array and returns its stats.
+            array_stats_fn: A function that takes an array and returns its stats,
+                or `None` if the array is not supported.
 
         Returns:
             A function that takes an array and returns its lovely representation.
@@ -74,9 +75,11 @@ def lovely_repr(dependencies: list[str]):
                     f"Missing dependencies: {', '.join(missing_deps)}. "
                     "Skipping lovely representation."
                 )
-                return repr(array)
+                return fallback_repr(array)
+
+            if (stats := array_stats_fn(array)) is None:
+                return fallback_repr(array)
 
-            stats = array_stats_fn(array)
             return format_tensor_stats(stats)
 
         return wrapper

nshutils/lovely/jax_.py

@@ -53,10 +53,14 @@ def _device(array: jax.Array) -> str:
     return f"{device.platform}:{device.id}"
 
 
-@lovely_repr(dependencies=["jax"])
-def jax_repr(array: jax.Array) -> LovelyStats:
+@lovely_repr(dependencies=["jax"], fallback_repr=jax.Array.__repr__)
+def jax_repr(array: jax.Array) -> LovelyStats | None:
     import jax.numpy as jnp
 
+    # For dtypes like `object` or `str`, we let the fallback repr handle it
+    if not jnp.issubdtype(array.dtype, jnp.number):
+        return None
+
     return {
         # Basic attributes
         "shape": array.shape,

nshutils/lovely/numpy_.py

@@ -51,8 +51,12 @@ def _dtype_str(array: np.ndarray) -> str:
     return dtype_base
 
 
-@lovely_repr(dependencies=["numpy"])
-def numpy_repr(array: np.ndarray) -> LovelyStats:
+@lovely_repr(dependencies=["numpy"], fallback_repr=np.array_repr)
+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):
+        return None
+
     return {
         # Basic attributes
         "shape": array.shape,

nshutils/lovely/torch_.py

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

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

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

{nshutils-0.22.4.dist-info → nshutils-0.22.6.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=QzB5E4kemxkTbz5oQrvGaAZl3n_KlU88aYXhVPetLuc,10458
+nshutils/actsave/_saver.py,sha256=GyoJpeJIIG_Y2Hr3MTXylTM6YNyoInf2mvKs0foU7no,10528
 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=c2XxNJlEdET2mP2gLzlYY1KHsEN4H9eDD_x8SptuBTA,4277
+nshutils/lovely/_base.py,sha256=SH8MY6wbucB7us4pU1cekoU2glanutEB8vDrLDkLEL0,4544
 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=mPH-tSOzWE27ymupllBnO4O6avvJJxpF2a1-G4dN5Ow,2214
-nshutils/lovely/numpy_.py,sha256=CI6pidaohK0tFw9mmSBTaDfjCKrFhY87ga0RIqTqSFk,3162
-nshutils/lovely/torch_.py,sha256=3wnXLa-1xwuQVk1fM50mBqOVDv5wZHjfnBwzOnjcFjg,2638
+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/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.4.dist-info/METADATA,sha256=eiBm_LbadsvvvXhKbHOyA6x5uRM8IO2CEHX9C3HtYww,4406
-nshutils-0.22.4.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-nshutils-0.22.4.dist-info/RECORD,,
+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,,