nshutils 0.22.4__tar.gz → 0.22.6__tar.gz

{nshutils-0.22.4 → nshutils-0.22.6}/PKG-INFO

@@ -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 → nshutils-0.22.6}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "nshutils"
-version = "0.22.4"
+version = "0.22.6"
 description = ""
 authors = [{ name = "Nima Shoghi", email = "nimashoghi@gmail.com" }]
 requires-python = ">=3.9,<4.0"

{nshutils-0.22.4 → nshutils-0.22.6}/src/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-0.22.4 → nshutils-0.22.6}/src/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-0.22.4 → nshutils-0.22.6}/src/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-0.22.4 → nshutils-0.22.6}/src/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-0.22.4 → nshutils-0.22.6}/src/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,