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

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.5.dist-info → nshutils-0.22.6.dist-info}/METADATA

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

{nshutils-0.22.5.dist-info → nshutils-0.22.6.dist-info}/RECORD

@@ -7,16 +7,16 @@ 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.5.dist-info/METADATA,sha256=dycLWOhxIaX5_uqvm-tIKFdftN0CexXMvPrvhsW-Vaw,4406
-nshutils-0.22.5.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-nshutils-0.22.5.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,,