nshutils 0.19.1__py3-none-any.whl → 0.21.0__py3-none-any.whl

nshutils/__init__.py

@@ -1,18 +1,9 @@
 from __future__ import annotations
 
-from . import actsave as actsave
-from . import typecheck as typecheck
-from .actsave import ActLoad as ActLoad
-from .actsave import ActSave as ActSave
-from .collections import apply_to_collection as apply_to_collection
-from .display import display as display
-from .logging import init_python_logging as init_python_logging
-from .logging import lovely as lovely
-from .logging import pretty as pretty
-from .snoop import snoop as snoop
-from .typecheck import tassert as tassert
-from .typecheck import typecheck_modules as typecheck_modules
-from .typecheck import typecheck_this_module as typecheck_this_module
+import lazy_loader as lazy
+
+__getattr__, __dir__, __all__ = lazy.attach_stub(__name__, __file__)
+
 
 try:
     from importlib.metadata import PackageNotFoundError, version

nshutils/__init__.pyi

@@ -0,0 +1,13 @@
+from . import actsave as actsave
+from . import lovely as lovely
+from . import typecheck as typecheck
+from .actsave import ActLoad as ActLoad
+from .actsave import ActSave as ActSave
+from .collections import apply_to_collection as apply_to_collection
+from .display import display as display
+from .logging import init_python_logging as init_python_logging
+from .logging import setup_logging as setup_logging
+from .snoop import snoop as snoop
+from .typecheck import tassert as tassert
+from .typecheck import typecheck_modules as typecheck_modules
+from .typecheck import typecheck_this_module as typecheck_this_module

nshutils/actsave/_saver.py

@@ -9,10 +9,10 @@ from dataclasses import dataclass
 from functools import wraps
 from logging import getLogger
 from pathlib import Path
-from typing import TYPE_CHECKING, Generic, TypeAlias, cast, overload
+from typing import TYPE_CHECKING, Generic, Literal, cast, overload
 
 import numpy as np
-from typing_extensions import Never, ParamSpec, TypeVar, override
+from typing_extensions import Never, ParamSpec, TypeAliasType, TypeVar, override
 from uuid_extensions import uuid7str
 
 from ..collections import apply_to_collection
@@ -21,25 +21,29 @@ if not TYPE_CHECKING:
     try:
         import torch  # type: ignore
 
-        Tensor: TypeAlias = torch.Tensor
+        Tensor = torch.Tensor
+        _torch_installed = True
     except ImportError:
         torch = None
+        _torch_installed = False
 
-        Tensor: TypeAlias = Never
+        Tensor = Never
 else:
     import torch  # type: ignore
 
-    Tensor: TypeAlias = torch.Tensor
-
+    Tensor = torch.Tensor
+    _torch_installed: Literal[True] = True
 
 log = getLogger(__name__)
 
-Value: TypeAlias = int | float | complex | bool | str | np.ndarray | Tensor | None
-ValueOrLambda: TypeAlias = Value | Callable[..., Value]
+Value = TypeAliasType(
+    "Value", int | float | complex | bool | str | np.ndarray | Tensor | None
+)
+ValueOrLambda = TypeAliasType("ValueOrLambda", Value | Callable[..., Value])
 
 
 def _torch_is_scripting() -> bool:
-    if torch is None:
+    if _torch_installed:
         return False
 
     return torch.jit.is_scripting()
@@ -54,7 +58,7 @@ def _to_numpy(activation: Value) -> np.ndarray:
         return np.array(activation)
     elif isinstance(activation, np.ndarray):
         return activation
-    elif isinstance(activation, Tensor):
+    elif _torch_installed and isinstance(activation, Tensor):
         activation_ = activation.detach()
         if activation_.is_floating_point():
             # NOTE: We need to convert to float32 because [b]float16 is not supported by numpy

nshutils/logging.py

@@ -2,12 +2,15 @@ from __future__ import annotations
 
 import logging
 from pathlib import Path
+from typing import TYPE_CHECKING
 
+if TYPE_CHECKING:
+    from .lovely._monkey_patch_all import Library
 
-def init_python_logging(
+
+def setup_logging(
     *,
-    lovely_tensors: bool = False,
-    lovely_numpy: bool = False,
+    lovely: bool | list[Library] = False,
     treescope: bool = False,
     treescope_autovisualize_arrays: bool = False,
     rich: bool = False,
@@ -15,25 +18,10 @@ def init_python_logging(
     log_level: int | str | None = logging.INFO,
     log_save_dir: Path | None = None,
 ):
-    if lovely_tensors:
-        try:
-            import lovely_tensors as _lovely_tensors  # type: ignore
-
-            _lovely_tensors.monkey_patch()
-        except ImportError:
-            logging.info(
-                "Failed to import `lovely_tensors`. Ignoring pretty PyTorch tensor formatting"
-            )
-
-    if lovely_numpy:
-        try:
-            import lovely_numpy as _lovely_numpy  # type: ignore
+    if lovely:
+        from .lovely._monkey_patch_all import monkey_patch
 
-            _lovely_numpy.set_config(repr=_lovely_numpy.lovely)
-        except ImportError:
-            logging.info(
-                "Failed to import `lovely_numpy`. Ignoring pretty numpy array formatting"
-            )
+        monkey_patch("auto" if lovely is True else lovely)
 
     if treescope:
         try:
@@ -77,49 +65,11 @@ def init_python_logging(
         datefmt="[%X]",
         handlers=log_handlers,
     )
-
-
-def pretty(
-    *,
-    lovely_tensors: bool = True,
-    lovely_numpy: bool = True,
-    treescope: bool = False,
-    treescope_autovisualize_arrays: bool = False,
-    log_level: int | str | None = logging.INFO,
-    log_save_dir: Path | None = None,
-    rich_log_handler: bool = False,
-    rich_tracebacks: bool = False,
-):
-    init_python_logging(
-        lovely_tensors=lovely_tensors,
-        lovely_numpy=lovely_numpy,
-        treescope=treescope,
-        treescope_autovisualize_arrays=treescope_autovisualize_arrays,
-        rich=rich_log_handler,
-        log_level=log_level,
-        log_save_dir=log_save_dir,
-        rich_tracebacks=rich_tracebacks,
+    logging.info(
+        "Logging initialized. "
+        f"Lovely: {lovely}, Treescope: {treescope}, Rich: {rich}, "
+        f"Log level: {log_level}, Log save dir: {log_save_dir}"
     )
 
 
-def lovely(
-    *,
-    lovely_tensors: bool = True,
-    lovely_numpy: bool = True,
-    treescope: bool = False,
-    treescope_autovisualize_arrays: bool = False,
-    log_level: int | str | None = logging.INFO,
-    log_save_dir: Path | None = None,
-    rich_log_handler: bool = False,
-    rich_tracebacks: bool = False,
-):
-    pretty(
-        lovely_tensors=lovely_tensors,
-        lovely_numpy=lovely_numpy,
-        treescope=treescope,
-        treescope_autovisualize_arrays=treescope_autovisualize_arrays,
-        log_level=log_level,
-        log_save_dir=log_save_dir,
-        rich_log_handler=rich_log_handler,
-        rich_tracebacks=rich_tracebacks,
-    )
+init_python_logging = setup_logging

nshutils/lovely/__init__.py

@@ -0,0 +1,10 @@
+from __future__ import annotations
+
+from ._monkey_patch_all import monkey_patch as monkey_patch
+from .config import LovelyConfig as LovelyConfig
+from .jax_ import jax_monkey_patch as jax_monkey_patch
+from .jax_ import jax_repr as jax_repr
+from .numpy_ import numpy_monkey_patch as numpy_monkey_patch
+from .numpy_ import numpy_repr as numpy_repr
+from .torch_ import torch_monkey_patch as torch_monkey_patch
+from .torch_ import torch_repr as torch_repr

nshutils/lovely/_base.py

@@ -0,0 +1,155 @@
+from __future__ import annotations
+
+import functools
+import importlib.util
+import logging
+from collections.abc import Callable, Iterator
+
+from typing_extensions import ParamSpec, TypeAliasType, TypeVar
+
+from ..util import ContextResource, resource_factory_contextmanager
+from .utils import LovelyStats, format_tensor_stats
+
+log = logging.getLogger(__name__)
+
+TArray = TypeVar("TArray", infer_variance=True)
+P = ParamSpec("P")
+
+LovelyStatsFn = TypeAliasType(
+    "LovelyStatsFn",
+    Callable[[TArray], LovelyStats],
+    type_params=(TArray,),
+)
+LovelyReprFn = TypeAliasType(
+    "LovelyReprFn",
+    Callable[[TArray], str],
+    type_params=(TArray,),
+)
+
+
+def _find_missing_deps(dependencies: list[str]):
+    missing_deps: list[str] = []
+
+    for dep in dependencies:
+        if importlib.util.find_spec(dep) is not None:
+            continue
+
+        missing_deps.append(dep)
+
+    return missing_deps
+
+
+def lovely_repr(dependencies: list[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.
+
+    Returns:
+        A decorator function that takes a function and returns a lovely representation function.
+
+    Example:
+        @lovely_repr(dependencies=["torch"])
+        def my_array_stats(array):
+            return {...}
+    """
+
+    def decorator_fn(array_stats_fn: LovelyStatsFn[TArray]) -> LovelyReprFn[TArray]:
+        """
+        Decorator to create a lovely representation function for an array.
+
+        Args:
+            array_stats_fn: A function that takes an array and returns its stats.
+
+        Returns:
+            A function that takes an array and returns its lovely representation.
+        """
+
+        @functools.wraps(array_stats_fn)
+        def wrapper(array: TArray) -> str:
+            if missing_deps := _find_missing_deps(dependencies):
+                log.warning(
+                    f"Missing dependencies: {', '.join(missing_deps)}. "
+                    "Skipping lovely representation."
+                )
+                return repr(array)
+
+            stats = array_stats_fn(array)
+            return format_tensor_stats(stats)
+
+        return wrapper
+
+    return decorator_fn
+
+
+LovelyMonkeyPatchInputFn = TypeAliasType(
+    "LovelyMonkeyPatchInputFn",
+    Callable[P, Iterator[None]],
+    type_params=(P,),
+)
+LovelyMonkeyPatchFn = TypeAliasType(
+    "LovelyMonkeyPatchFn",
+    Callable[P, ContextResource[None]],
+    type_params=(P,),
+)
+
+
+def _nullcontext_generator():
+    """A generator that does nothing."""
+    yield
+
+
+def _wrap_monkey_patch_fn(
+    monkey_patch_fn: LovelyMonkeyPatchInputFn[P],
+    dependencies: list[str],
+) -> LovelyMonkeyPatchInputFn[P]:
+    @functools.wraps(monkey_patch_fn)
+    def wrapper(*args: P.args, **kwargs: P.kwargs) -> Iterator[None]:
+        if missing_deps := _find_missing_deps(dependencies):
+            log.warning(
+                f"Missing dependencies: {', '.join(missing_deps)}. "
+                "Skipping monkey patch."
+            )
+            return _nullcontext_generator()
+
+        return monkey_patch_fn(*args, **kwargs)
+
+    return wrapper
+
+
+def monkey_patch_contextmanager(dependencies: list[str]):
+    """
+    Decorator to create a monkey patch 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.
+
+    Returns:
+        A decorator function that takes a function and returns a monkey patch function.
+
+    Example:
+        @monkey_patch_contextmanager(dependencies=["torch"])
+        def my_array_monkey_patch():
+            ...
+    """
+
+    def decorator_fn(
+        monkey_patch_fn: LovelyMonkeyPatchInputFn[P],
+    ) -> LovelyMonkeyPatchFn[P]:
+        """
+        Decorator to create a monkey patch function for an array.
+
+        Args:
+            monkey_patch_fn: A function that applies the monkey patch.
+
+        Returns:
+            A function that applies the monkey patch.
+        """
+
+        wrapped_fn = _wrap_monkey_patch_fn(monkey_patch_fn, dependencies)
+        return resource_factory_contextmanager(wrapped_fn)
+
+    return decorator_fn

nshutils/lovely/_monkey_patch_all.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+import contextlib
+import importlib.util
+import logging
+from typing import Literal
+
+from typing_extensions import TypeAliasType, assert_never
+
+from ..util import resource_factory_contextmanager
+
+Library = TypeAliasType("Library", Literal["numpy", "torch", "jax"])
+
+log = logging.getLogger(__name__)
+
+
+def _find_deps() -> list[Library]:
+    """
+    Find available libraries for monkey patching.
+    """
+    deps: list[Library] = []
+    if importlib.util.find_spec("torch") is not None:
+        deps.append("torch")
+    if importlib.util.find_spec("jax") is not None:
+        deps.append("jax")
+    if importlib.util.find_spec("numpy") is not None:
+        deps.append("numpy")
+    return deps
+
+
+@resource_factory_contextmanager
+def monkey_patch(libraries: list[Library] | Literal["auto"] = "auto"):
+    if libraries == "auto":
+        libraries = _find_deps()
+
+    if not libraries:
+        raise ValueError(
+            "No libraries found for monkey patching. "
+            "Please install numpy, torch, or jax."
+        )
+
+    with contextlib.ExitStack() as stack:
+        for library in libraries:
+            match library:
+                case "torch":
+                    from .torch_ import torch_monkey_patch
+
+                    stack.enter_context(torch_monkey_patch())
+                case "jax":
+                    from .jax_ import jax_monkey_patch
+
+                    stack.enter_context(jax_monkey_patch())
+                case "numpy":
+                    from .numpy_ import numpy_monkey_patch
+
+                    stack.enter_context(numpy_monkey_patch())
+                case _:
+                    assert_never(library)
+
+        log.info(
+            f"Monkey patched libraries: {', '.join(libraries)}. "
+            "You can now use the lovely functions with these libraries."
+        )
+        yield
+        log.info(
+            f"Unmonkey patched libraries: {', '.join(libraries)}. "
+            "You can now use the lovely functions with these libraries."
+        )

nshutils/lovely/config.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+from typing import ClassVar
+
+import nshconfig as C
+from typing_extensions import Self
+
+
+class LovelyConfig(C.Config):
+    """
+    This class is used to manage the configuration of the Lovely library.
+    It inherits from the Config class in the nshconfig module.
+    """
+
+    precision: int = 3
+    """Number of digits after the decimal point."""
+
+    threshold_max: int = 3
+    """Absolute values larger than 10^3 use scientific notation."""
+
+    threshold_min: int = -4
+    """Absolute values smaller than 10^-4 use scientific notation."""
+
+    sci_mode: bool | None = None
+    """Force scientific notation (None=auto)."""
+
+    show_mem_above: int = 1024
+    """Show memory size if above this threshold (bytes)."""
+
+    color: bool = True
+    """Use ANSI colors in text."""
+
+    indent: int = 2
+    """Indent for nested representation."""
+
+    config_instance: ClassVar[Self | None] = None
+    """Singleton instance of the LovelyConfig class."""
+
+    @classmethod
+    def instance(cls) -> Self:
+        """
+        Get the singleton instance of the LovelyConfig class.
+        If it doesn't exist, create it.
+        """
+        if cls.config_instance is None:
+            cls.config_instance = cls()
+        return cls.config_instance

nshutils/lovely/jax_.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, cast
+
+import numpy as np
+
+from ._base import lovely_repr, monkey_patch_contextmanager
+from .utils import LovelyStats, array_stats, patch_to
+
+if TYPE_CHECKING:
+    import jax
+
+
+def _type_name(array: jax.Array):
+    type_name = type(array).__name__.rsplit(".", 1)[-1]
+    return "array" if type_name == "ArrayImpl" else type_name
+
+
+_DT_NAMES = {
+    "float16": "f16",
+    "float32": "f32",
+    "float64": "f64",
+    "uint8": "u8",
+    "uint16": "u16",
+    "uint32": "u32",
+    "uint64": "u64",
+    "int8": "i8",
+    "int16": "i16",
+    "int32": "i32",
+    "int64": "i64",
+    "bfloat16": "bf16",
+    "complex64": "c64",
+    "complex128": "c128",
+}
+
+
+def _dtype_str(array: jax.Array) -> str:
+    dtype_base = str(array.dtype).rsplit(".", 1)[-1]
+    dtype_base = _DT_NAMES.get(dtype_base, dtype_base)
+    return dtype_base
+
+
+def _device(array: jax.Array) -> str:
+    from jaxlib.xla_extension import Device
+
+    if callable(device := array.device):
+        device = device()
+
+    device = cast(Device, device)
+    if device.platform == "cpu":
+        return "cpu"
+
+    return f"{device.platform}:{device.id}"
+
+
+@lovely_repr(dependencies=["jax"])
+def jax_repr(array: jax.Array) -> LovelyStats:
+    import jax.numpy as jnp
+
+    return {
+        # Basic attributes
+        "shape": array.shape,
+        "size": array.size,
+        "nbytes": array.nbytes,
+        "type_name": _type_name(array),
+        # Dtype
+        "dtype_str": _dtype_str(array),
+        "is_complex": jnp.iscomplexobj(array),
+        # Device
+        "device": _device(array),
+        # Depending of whether the tensor is complex or not, we will call the appropriate stats function
+        **array_stats(np.asarray(array)),
+    }
+
+
+@monkey_patch_contextmanager(dependencies=["jax"])
+def jax_monkey_patch():
+    from jax._src import array
+
+    prev_repr = array.ArrayImpl.__repr__
+    prev_str = array.ArrayImpl.__str__
+    try:
+        patch_to(array.ArrayImpl, "__repr__", jax_repr)
+        patch_to(array.ArrayImpl, "__str__", jax_repr)
+
+        yield
+    finally:
+        patch_to(array.ArrayImpl, "__repr__", prev_repr)
+        patch_to(array.ArrayImpl, "__str__", prev_str)

nshutils/lovely/numpy_.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+
+import logging
+
+import numpy as np
+
+from ._base import lovely_repr, monkey_patch_contextmanager
+from .utils import LovelyStats, array_stats
+
+
+def _type_name(array: np.ndarray):
+    return (
+        "array"
+        if type(array) is np.ndarray
+        else type(array).__name__.rsplit(".", 1)[-1]
+    )
+
+
+_DT_NAMES = {
+    "float16": "f16",
+    "float32": "f32",
+    "float64": "",  # Default dtype in numpy
+    "uint8": "u8",
+    "uint16": "u16",
+    "uint32": "u32",
+    "uint64": "u64",
+    "int8": "i8",
+    "int16": "i16",
+    "int32": "i32",
+    "int64": "i64",
+    "complex64": "c64",
+    "complex128": "c128",
+}
+
+
+def _dtype_str(array: np.ndarray) -> str:
+    dtype_base = str(array.dtype).rsplit(".", 1)[-1]
+    dtype_base = _DT_NAMES.get(dtype_base, dtype_base)
+    return dtype_base
+
+
+@lovely_repr(dependencies=["numpy"])
+def numpy_repr(array: np.ndarray) -> LovelyStats:
+    return {
+        # Basic attributes
+        "shape": array.shape,
+        "size": array.size,
+        "nbytes": array.nbytes,
+        "type_name": _type_name(array),
+        # Dtype
+        "dtype_str": _dtype_str(array),
+        "is_complex": np.iscomplexobj(array),
+        # Depending of whether the tensor is complex or not, we will call the appropriate stats function
+        **array_stats(array),
+    }
+
+
+@monkey_patch_contextmanager(dependencies=["numpy"])
+def numpy_monkey_patch():
+    try:
+        np.set_printoptions(override_repr=numpy_repr)
+        logging.info(
+            f"Numpy monkey patching: using {numpy_repr.__name__} for numpy arrays. "
+            f"{np.get_printoptions()=}"
+        )
+        yield
+    finally:
+        np.set_printoptions(override_repr=None)
+        logging.info(
+            f"Numpy unmonkey patching: using {numpy_repr.__name__} for numpy arrays. "
+            f"{np.get_printoptions()=}"
+        )

nshutils/lovely/torch_.py

@@ -0,0 +1,99 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+from ._base import lovely_repr, monkey_patch_contextmanager
+from .utils import LovelyStats, array_stats, patch_to
+
+if TYPE_CHECKING:
+    import torch
+
+
+def _type_name(tensor: torch.Tensor):
+    import torch
+
+    return (
+        "tensor"
+        if type(tensor) is torch.Tensor
+        else type(tensor).__name__.split(".")[-1]
+    )
+
+
+_DT_NAMES = {
+    "float32": "",  # Default dtype
+    "float16": "f16",
+    "float64": "f64",
+    "bfloat16": "bf16",
+    "uint8": "u8",
+    "int8": "i8",
+    "int16": "i16",
+    "int32": "i32",
+    "int64": "i64",
+    "complex32": "c32",
+    "complex64": "c64",
+    "complex128": "c128",
+}
+
+
+def _dtype_str(tensor: torch.Tensor) -> str:
+    dtype_base = str(tensor.dtype).rsplit(".", 1)[-1]
+    dtype_base = _DT_NAMES.get(dtype_base, dtype_base)
+    return dtype_base
+
+
+def _to_np(tensor: torch.Tensor) -> np.ndarray:
+    import torch
+
+    # Get tensor data as CPU NumPy array for analysis
+    t_cpu = tensor.detach().cpu()
+
+    # Convert bfloat16 to float32 for numpy compatibility
+    if tensor.dtype == torch.bfloat16:
+        t_cpu = t_cpu.to(torch.float32)
+
+    # Convert to NumPy
+    t_np = t_cpu.numpy()
+
+    return t_np
+
+
+@lovely_repr(dependencies=["torch"])
+def torch_repr(tensor: torch.Tensor) -> LovelyStats:
+    return {
+        # Basic attributes
+        "shape": tensor.shape,
+        "size": tensor.numel(),
+        "nbytes": tensor.element_size() * tensor.numel(),
+        "type_name": _type_name(tensor),
+        # Device
+        "device": str(tensor.device) if tensor.device else None,
+        "is_meta": device.type == "meta" if (device := tensor.device) else False,
+        # Grad
+        "requires_grad": tensor.requires_grad,
+        # Dtype
+        "dtype_str": _dtype_str(tensor),
+        "is_complex": tensor.is_complex(),
+        # Depending of whether the tensor is complex or not, we will call the appropriate stats function
+        **array_stats(_to_np(tensor)),
+    }
+
+
+@monkey_patch_contextmanager(dependencies=["torch"])
+def torch_monkey_patch():
+    import torch
+
+    original_repr = torch.Tensor.__repr__
+    original_str = torch.Tensor.__str__
+    original_parameter_repr = torch.nn.Parameter.__repr__
+    try:
+        patch_to(torch.Tensor, "__repr__", torch_repr)
+        patch_to(torch.Tensor, "__str__", torch_repr)
+        del torch.nn.Parameter.__repr__
+
+        yield
+    finally:
+        patch_to(torch.Tensor, "__repr__", original_repr)
+        patch_to(torch.Tensor, "__str__", original_str)
+        patch_to(torch.nn.Parameter, "__repr__", original_parameter_repr)

nshutils/lovely/utils.py

@@ -0,0 +1,345 @@
+from __future__ import annotations
+
+import sys
+from collections import defaultdict
+from collections.abc import Callable, Sequence
+from typing import Any
+
+import numpy as np
+from typing_extensions import TypedDict
+
+from .config import LovelyConfig
+
+
+class LovelyStats(TypedDict, total=False):
+    """Statistics for tensor representation."""
+
+    # Basic tensor information
+    shape: Sequence[int]
+    size: int
+    type_name: str
+    dtype_str: str
+    device: str | None
+    nbytes: int
+    requires_grad: bool
+    is_meta: bool
+
+    # Content flags
+    all_zeros: bool
+    has_nan: bool
+    has_pos_inf: bool
+    has_neg_inf: bool
+    is_complex: bool
+
+    # Numeric statistics
+    min: float | None
+    max: float | None
+    mean: float | None
+    std: float | None
+
+    # Complex number statistics
+    mag_min: float | None
+    mag_max: float | None
+    real_min: float | None
+    real_max: float | None
+    imag_min: float | None
+    imag_max: float | None
+
+    # Representation
+    values_str: str | None
+
+
+# Formatting utilities
+def sci_mode(f: float) -> bool:
+    """Determine if a float should be displayed in scientific notation."""
+    config = LovelyConfig.instance()
+    return (abs(f) < 10**config.threshold_min) or (abs(f) > 10**config.threshold_max)
+
+
+def pretty_str(x: Any) -> str:
+    """Format a number or array for pretty display.
+
+    Works with scalars, numpy arrays, torch tensors, and jax arrays.
+    """
+    if isinstance(x, int):
+        return f"{x}"
+    elif isinstance(x, float):
+        if x == 0.0:
+            return "0."
+
+        sci = (
+            sci_mode(x)
+            if LovelyConfig.instance().sci_mode is None
+            else LovelyConfig.instance().sci_mode
+        )
+        fmt = f"{{:.{LovelyConfig.instance().precision}{'e' if sci else 'f'}}}"
+        return fmt.format(x)
+    elif isinstance(x, complex):
+        # Handle complex numbers
+        real_part = pretty_str(x.real)
+        imag_part = pretty_str(abs(x.imag))
+        sign = "+" if x.imag >= 0 else "-"
+        return f"{real_part}{sign}{imag_part}j"
+
+    # Handle array-like objects
+    try:
+        if hasattr(x, "ndim") and x.ndim == 0:
+            return pretty_str(x.item())
+        elif hasattr(x, "shape") and len(x.shape) > 0:
+            slices = [pretty_str(x[i]) for i in range(min(x.shape[0], 10))]
+            if x.shape[0] > 10:
+                slices.append("...")
+            return "[" + ", ".join(slices) + "]"
+    except:
+        pass
+
+    # Fallback
+    return str(x)
+
+
+def sparse_join(items: list[str | None], sep: str = " ") -> str:
+    """Join non-empty strings with a separator."""
+    return sep.join([item for item in items if item])
+
+
+def ansi_color(s: str, col: str, use_color: bool = True) -> str:
+    """Add ANSI color to a string if use_color is True."""
+    if not use_color:
+        return s
+
+    style = defaultdict(str)
+    style["grey"] = "\x1b[38;2;127;127;127m"
+    style["red"] = "\x1b[31m"
+    end_style = "\x1b[0m"
+
+    return style[col] + s + end_style
+
+
+def bytes_to_human(num_bytes: int) -> str:
+    """Convert bytes to a human-readable string (b, Kb, Mb, Gb)."""
+    units = ["b", "Kb", "Mb", "Gb"]
+
+    value = num_bytes
+    matched_unit: str | None = None
+    for unit in units:
+        if value < 1024 / 10:
+            matched_unit = unit
+            break
+        value /= 1024.0
+
+    assert matched_unit is not None, "No matching unit found"
+
+    if value % 1 == 0 or value >= 10:
+        return f"{round(value)}{matched_unit}"
+    else:
+        return f"{value:.1f}{matched_unit}"
+
+
+def in_debugger() -> bool:
+    """Returns True if running in a debugger."""
+    return getattr(sys, "gettrace", None) is not None and sys.gettrace() is not None
+
+
+# Common tensor representation
+def format_tensor_stats(tensor_stats: LovelyStats, color: bool | None = None) -> str:
+    """Format tensor stats into a pretty string representation."""
+    conf = LovelyConfig.instance()
+    if color is None:
+        color = conf.color
+    if in_debugger():
+        color = False
+
+    # Basic tensor info
+    shape_str = str(list(shape)) if (shape := tensor_stats.get("shape")) else None
+    type_str = sparse_join([tensor_stats.get("type_name"), shape_str], sep="")
+
+    # Calculate memory usage
+    numel = None
+    if (size := tensor_stats.get("size")) and (nbytes := tensor_stats.get("nbytes")):
+        shape = tensor_stats.get("shape", [])
+
+        if shape and max(shape) != size:
+            numel = f"n={size}"
+            if conf.show_mem_above <= nbytes:
+                numel = sparse_join([numel, f"({bytes_to_human(nbytes)})"])
+        elif conf.show_mem_above <= nbytes:
+            numel = bytes_to_human(nbytes)
+
+    # Handle empty tensors
+    if tensor_stats.get("size", 0) == 0:
+        common = ansi_color("empty", "grey", color)
+    # Handle all zeros
+    elif tensor_stats.get("all_zeros"):
+        common = ansi_color("all_zeros", "grey", color)
+    # Handle complex tensors
+    elif tensor_stats.get("is_complex"):
+        complex_info = []
+
+        # For magnitude stats
+        if (mag_min := tensor_stats.get("mag_min")) is not None and (
+            mag_max := tensor_stats.get("mag_max")
+        ) is not None:
+            complex_info.append(f"|z|∈[{pretty_str(mag_min)}, {pretty_str(mag_max)}]")
+
+        # For real part stats
+        if (real_min := tensor_stats.get("real_min")) is not None and (
+            real_max := tensor_stats.get("real_max")
+        ) is not None:
+            complex_info.append(f"Re∈[{pretty_str(real_min)}, {pretty_str(real_max)}]")
+
+        # For imaginary part stats
+        if (imag_min := tensor_stats.get("imag_min")) is not None and (
+            imag_max := tensor_stats.get("imag_max")
+        ) is not None:
+            complex_info.append(f"Im∈[{pretty_str(imag_min)}, {pretty_str(imag_max)}]")
+
+        common = sparse_join(complex_info)
+    # Handle normal tensors with stats
+    elif (min_val := tensor_stats.get("min")) is not None and (
+        max_val := tensor_stats.get("max")
+    ) is not None:
+        minmax = None
+        meanstd = None
+
+        if tensor_stats.get("size", 0) > 2:
+            minmax = f"x∈[{pretty_str(min_val)}, {pretty_str(max_val)}]"
+
+        if (
+            (mean := tensor_stats.get("mean")) is not None
+            and (std := tensor_stats.get("std")) is not None
+            and tensor_stats.get("size", 0) >= 2
+        ):
+            meanstd = f"μ={pretty_str(mean)} σ={pretty_str(std)}"
+
+        common = sparse_join([minmax, meanstd])
+    else:
+        common = None
+
+    # Handle warnings
+    warnings = []
+    if tensor_stats.get("has_nan"):
+        warnings.append(ansi_color("NaN!", "red", color))
+    if tensor_stats.get("has_pos_inf"):
+        warnings.append(ansi_color("+Inf!", "red", color))
+    if tensor_stats.get("has_neg_inf"):
+        warnings.append(ansi_color("-Inf!", "red", color))
+
+    attention = sparse_join(warnings)
+    common = sparse_join([common, attention])
+
+    # Other tensor attributes
+    dtype = tensor_stats.get("dtype_str", "")
+    device = tensor_stats.get("device")
+    grad = "grad" if tensor_stats.get("requires_grad") else None
+
+    # Format values for small tensors
+    vals = None
+    if (
+        0 < tensor_stats.get("size", 0) <= 10
+        and tensor_stats.get("is_meta", False) is False
+    ):
+        vals = tensor_stats.get("values_str")
+
+    # Join all parts
+    result = sparse_join([type_str, dtype, numel, common, grad, device, vals])
+
+    return result
+
+
+def real_stats(array: np.ndarray) -> LovelyStats:
+    stats: LovelyStats = {}
+
+    # Check for special values
+    stats["has_nan"] = bool(np.isnan(array).any())
+    stats["has_pos_inf"] = bool(np.isposinf(array).any())
+    stats["has_neg_inf"] = bool(np.isneginf(array).any())
+
+    # Only compute min/max/mean/std for good data
+    good_data = array[np.isfinite(array)]
+
+    if len(good_data) > 0:
+        stats["min"] = float(good_data.min())
+        stats["max"] = float(good_data.max())
+        stats["all_zeros"] = stats["min"] == 0 and stats["max"] == 0 and array.size > 1
+
+        if len(good_data) > 1:
+            stats["mean"] = float(good_data.mean())
+            stats["std"] = float(good_data.std())
+
+    # Get string representation of values for small tensors
+    if 0 < array.size <= 10:
+        stats["values_str"] = pretty_str(array)
+
+    return stats
+
+
+def complex_stats(array: np.ndarray) -> LovelyStats:
+    stats: LovelyStats = {}
+
+    # Calculate magnitude (absolute value)
+    magnitude = np.abs(array)
+
+    # Check for special values in real or imaginary parts
+    stats["has_nan"] = bool(np.isnan(array.real).any() or np.isnan(array.imag).any())
+
+    # Get statistics for magnitude
+    good_mag = magnitude[np.isfinite(magnitude)]
+    if len(good_mag) > 0:
+        stats["mag_min"] = float(good_mag.min())
+        stats["mag_max"] = float(good_mag.max())
+        stats["all_zeros"] = (
+            stats["mag_min"] == 0 and stats["mag_max"] == 0 and array.size > 1
+        )
+
+    # Get statistics for real and imaginary parts
+    real_part = array.real
+    imag_part = array.imag
+
+    good_real = real_part[np.isfinite(real_part)]
+    good_imag = imag_part[np.isfinite(imag_part)]
+
+    if len(good_real := real_part[np.isfinite(real_part)]):
+        stats["real_min"] = float(good_real.min())
+        stats["real_max"] = float(good_real.max())
+
+    if len(good_imag := imag_part[np.isfinite(imag_part)]):
+        stats["imag_min"] = float(good_imag.min())
+        stats["imag_max"] = float(good_imag.max())
+
+    # Get string representation of values for small tensors
+    if 0 < array.size <= 10:
+        stats["values_str"] = pretty_str(array)
+
+    return stats
+
+
+def array_stats(array: np.ndarray, ignore_empty: bool = True) -> LovelyStats:
+    """Compute all statistics for a given array.
+
+    Args:
+        array (np.ndarray): The input array.
+        ignore_empty (bool): If True, ignore empty arrays.
+
+    Returns:
+        LovelyStats: A dictionary containing the computed statistics.
+    """
+    if ignore_empty and array.size == 0:
+        return {}
+
+    if np.iscomplexobj(array):
+        return complex_stats(array)
+    else:
+        return real_stats(array)
+
+
+def patch_to(
+    cls: type,
+    name: str,
+    func: Callable[[Any], Any],
+    as_property: bool = False,
+) -> None:
+    """Simple patch_to implementation to avoid fastcore dependency."""
+    if as_property:
+        setattr(cls, name, property(func))
+    else:
+        setattr(cls, name, func)

nshutils/typecheck.py

@@ -1,9 +1,9 @@
 from __future__ import annotations
 
+import logging
 import os
 import sys
 from collections.abc import Sequence
-from logging import getLogger
 from types import FrameType as _FrameType
 from typing import Any
 
@@ -52,7 +52,8 @@ try:
     import jax  # type: ignore
 except ImportError:
     jax = None
-log = getLogger(__name__)
+
+log = logging.getLogger(__name__)
 
 DISABLE_ENV_KEY = "NSHUTILS_DISABLE_TYPECHECKING"
 

nshutils/util.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+import contextlib
+import functools
+from collections.abc import Callable, Iterator
+from typing import Any, Generic
+
+from typing_extensions import ParamSpec, TypeVar, override
+
+R = TypeVar("R")
+P = ParamSpec("P")
+
+
+class ContextResource(contextlib.AbstractContextManager[R], Generic[R]):
+    """A class that provides both direct access to a resource and context management."""
+
+    def __init__(self, resource: R, cleanup_func: Callable[[R], Any]):
+        self.resource = resource
+        self._cleanup_func = cleanup_func
+
+    @override
+    def __enter__(self) -> R:
+        """When used as a context manager, return the wrapped resource."""
+        return self.resource
+
+    @override
+    def __exit__(self, *exc_info) -> None:
+        """Clean up the resource when exiting the context."""
+        self._cleanup_func(self.resource)
+
+    def close(self) -> None:
+        """Explicitly clean up the resource."""
+        self._cleanup_func(self.resource)
+
+
+def resource_factory(
+    create_func: Callable[P, R], cleanup_func: Callable[[R], None]
+) -> Callable[P, ContextResource[R]]:
+    """
+    Create a factory function that returns a ContextResource.
+
+    Args:
+        create_func: Function that creates the resource
+        cleanup_func: Function that cleans up the resource
+
+    Returns:
+        A function that returns a ContextResource wrapping the created resource
+    """
+
+    @functools.wraps(create_func)
+    def wrapper(*args: P.args, **kwargs: P.kwargs) -> ContextResource[R]:
+        resource = create_func(*args, **kwargs)
+        return ContextResource(resource, cleanup_func)
+
+    return wrapper
+
+
+def resource_factory_from_context_fn(
+    context_func: Callable[P, contextlib.AbstractContextManager[R]],
+) -> Callable[P, ContextResource[R]]:
+    """
+    Create a factory function that returns a ContextResource.
+
+    Args:
+        context_func: Function that creates the resource
+
+    Returns:
+        A function that returns a ContextResource wrapping the created resource
+    """
+
+    @functools.wraps(context_func)
+    def wrapper(*args: P.args, **kwargs: P.kwargs) -> ContextResource[R]:
+        context = context_func(*args, **kwargs)
+        resource = context.__enter__()
+        return ContextResource(resource, lambda _: context.__exit__(None, None, None))
+
+    return wrapper
+
+
+def resource_factory_contextmanager(
+    context_func: Callable[P, Iterator[R]],
+) -> Callable[P, ContextResource[R]]:
+    """
+    Create a factory function that returns a ContextResource.
+
+    Args:
+        context_func: Generator function that creates the resource, yields it, and cleans up the resource when done.
+
+    Returns:
+        A function that returns a ContextResource wrapping the created resource
+    """
+    return resource_factory_from_context_fn(contextlib.contextmanager(context_func))

{nshutils-0.19.1.dist-info → nshutils-0.21.0.dist-info}/METADATA

@@ -1,24 +1,30 @@
 Metadata-Version: 2.3
 Name: nshutils
-Version: 0.19.1
+Version: 0.21.0
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com
-Requires-Python: >=3.10,<4.0
+Requires-Python: >=3.9,<4.0
 Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Provides-Extra: extra
+Provides-Extra: pprint
+Provides-Extra: snoop
 Requires-Dist: beartype
 Requires-Dist: jaxtyping
-Requires-Dist: lovely-numpy ; extra == "extra"
-Requires-Dist: lovely-tensors ; extra == "extra"
+Requires-Dist: lazy-loader
+Requires-Dist: nshconfig
 Requires-Dist: numpy
 Requires-Dist: pysnooper ; extra == "extra"
-Requires-Dist: rich ; extra == "extra"
-Requires-Dist: treescope ; extra == "extra"
+Requires-Dist: pysnooper ; extra == "snoop"
+Requires-Dist: rich[jupyter] ; extra == "extra"
+Requires-Dist: rich[jupyter] ; extra == "pprint"
+Requires-Dist: treescope ; (python_version >= "3.10") and (extra == "extra")
+Requires-Dist: treescope ; (python_version >= "3.10") and (extra == "pprint")
 Requires-Dist: typing-extensions
 Requires-Dist: uuid7
 Project-URL: Homepage, https://github.com/nimashoghi/nshutils

nshutils-0.21.0.dist-info/RECORD

@@ -0,0 +1,22 @@
+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=LulgC_B7oYtsGjW_I_pnSLrf1k9lmMvHqRxcmfqqrjU,10441
+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/_monkey_patch_all.py,sha256=WZsC6Xp5-Z2GBd6xyZZEsD4C2xNyZy0YBfjZzAy3m8M,2028
+nshutils/lovely/config.py,sha256=jsUK_kfEvthL94qxpHxM-Xobdv67sWZpnH_ag4HLTNo,1274
+nshutils/lovely/jax_.py,sha256=mPH-tSOzWE27ymupllBnO4O6avvJJxpF2a1-G4dN5Ow,2214
+nshutils/lovely/numpy_.py,sha256=iHA4BCJIW9IU6DXKEfYbh9RA2xyeXeL0tyHdEja89Sw,1866
+nshutils/lovely/torch_.py,sha256=3wnXLa-1xwuQVk1fM50mBqOVDv5wZHjfnBwzOnjcFjg,2638
+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.21.0.dist-info/METADATA,sha256=Qa4bVvnzN2Ip9fbmuaToECZbRSKhgj-zmaNqTKmXXJ4,4431
+nshutils-0.21.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+nshutils-0.21.0.dist-info/RECORD,,

{nshutils-0.19.1.dist-info → nshutils-0.21.0.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: poetry-core 2.1.2
+Generator: poetry-core 2.1.3
 Root-Is-Purelib: true
 Tag: py3-none-any

nshutils-0.19.1.dist-info/RECORD

@@ -1,12 +0,0 @@
-nshutils/__init__.py,sha256=9RJO-Bt7uN-KHf7Yi9G7s_5AL8n-dwF2SClNMNQpiQE,985
-nshutils/actsave/__init__.py,sha256=hAVsog9d1g3_rQN1TRslrl6sK1PhCGbjy8PPUAmJI58,203
-nshutils/actsave/_loader.py,sha256=mof3HezUNvLliz7macstX6ewXW05L0Mtv3zJyrbmImg,4640
-nshutils/actsave/_saver.py,sha256=Kor7PEk__noRDEAfCZ_I4vxql5UfWSIQIQ1OmW2RRTI,10290
-nshutils/collections.py,sha256=QWGyANmo4Efq4XRNHDSTE9tRLStwEZHGwE0ATHR-Vqo,5233
-nshutils/display.py,sha256=Ge63yllx7gi-MKL3mKQeQ5doql_nj56-o5aoTVmusDg,1473
-nshutils/logging.py,sha256=-6IB0GTDDS8ue1H2tzkv_OLf4bZVN1ywL08TlDZWbtQ,3737
-nshutils/snoop.py,sha256=7d7_Q5sJmINL1J29wcnxEvpV95zvZYNoVn5frCq-rww,7393
-nshutils/typecheck.py,sha256=UOUYfa72wTmc-a7VQw52tKFb4U10xq1qcZuEWc2sAd8,5588
-nshutils-0.19.1.dist-info/METADATA,sha256=nxPB1rd9Lv_CnL03BAMWGvBnRXYZ4VLsFTYpyAlhSAY,4168
-nshutils-0.19.1.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-nshutils-0.19.1.dist-info/RECORD,,