nshutils 0.19.0__tar.gz → 0.20.0__tar.gz

{nshutils-0.19.0 → nshutils-0.20.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: nshutils
-Version: 0.19.0
+Version: 0.20.0
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com
@@ -11,14 +11,19 @@ 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: pysnooper ; extra == "snoop"
+Requires-Dist: rich[jupyter] ; extra == "extra"
+Requires-Dist: rich[jupyter] ; extra == "pprint"
 Requires-Dist: treescope ; extra == "extra"
+Requires-Dist: treescope ; extra == "pprint"
 Requires-Dist: typing-extensions
 Requires-Dist: uuid7
 Project-URL: Homepage, https://github.com/nimashoghi/nshutils

{nshutils-0.19.0 → nshutils-0.20.0}/pyproject.toml

@@ -1,15 +1,25 @@
 [project]
 name = "nshutils"
-version = "0.19.0"
+version = "0.20.0"
 description = ""
 authors = [{ name = "Nima Shoghi", email = "nimashoghi@gmail.com" }]
 requires-python = ">=3.10,<4.0"
 readme = "README.md"
 
-dependencies = ["numpy", "jaxtyping", "typing-extensions", "beartype", "uuid7"]
+dependencies = [
+    "lazy-loader",
+    "numpy",
+    "typing-extensions",
+    "jaxtyping",
+    "beartype",
+    "uuid7",
+    "nshconfig",
+]
 
 [project.optional-dependencies]
-extra = ["pysnooper", "lovely-numpy", "lovely-tensors", "rich", "treescope"]
+snoop = ["pysnooper"]
+pprint = ["rich[jupyter]", "treescope"]
+extra = ["pysnooper", "rich[jupyter]", "treescope"]
 
 [project.urls]
 homepage = "https://github.com/nimashoghi/nshutils"

nshutils-0.20.0/src/nshutils/__init__.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+import lazy_loader as lazy
+
+__getattr__, __dir__, __all__ = lazy.attach_stub(__name__, __file__)
+
+
+try:
+    from importlib.metadata import PackageNotFoundError, version
+except ImportError:
+    # For Python <3.8
+    from importlib_metadata import (  # pyright: ignore[reportMissingImports]
+        PackageNotFoundError,
+        version,
+    )
+
+try:
+    __version__ = version(__name__)
+except PackageNotFoundError:
+    __version__ = "unknown"

nshutils-0.19.0/src/nshutils/__init__.py → nshutils-0.20.0/src/nshutils/__init__.pyi

@@ -1,14 +1,12 @@
-from __future__ import annotations
-
 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 lovely as lovely
-from .logging import pretty as pretty
+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

{nshutils-0.19.0 → nshutils-0.20.0}/src/nshutils/actsave/_saver.py

@@ -12,7 +12,7 @@ from pathlib import Path
 from typing import TYPE_CHECKING, Generic, TypeAlias, 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,21 +21,23 @@ if not TYPE_CHECKING:
     try:
         import torch  # type: ignore
 
-        Tensor: TypeAlias = torch.Tensor
+        Tensor = torch.Tensor
     except ImportError:
         torch = None
 
-        Tensor: TypeAlias = Never
+        Tensor = Never
 else:
     import torch  # type: ignore
 
-    Tensor: TypeAlias = torch.Tensor
+    Tensor = torch.Tensor
 
 
 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:

nshutils-0.20.0/src/nshutils/logging.py

@@ -0,0 +1,75 @@
+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 setup_logging(
+    *,
+    lovely: bool | list[Library] = False,
+    treescope: bool = False,
+    treescope_autovisualize_arrays: bool = False,
+    rich: bool = False,
+    rich_tracebacks: bool = False,
+    log_level: int | str | None = logging.INFO,
+    log_save_dir: Path | None = None,
+):
+    if lovely:
+        from .lovely._monkey_patch_all import monkey_patch
+
+        monkey_patch("auto" if lovely is True else lovely)
+
+    if treescope:
+        try:
+            # Check if we're in a Jupyter environment
+            from IPython import get_ipython
+
+            if get_ipython() is not None:
+                import treescope as _treescope  # type: ignore
+
+                _treescope.basic_interactive_setup(
+                    autovisualize_arrays=treescope_autovisualize_arrays
+                )
+            else:
+                logging.info(
+                    "Treescope setup is only supported in Jupyter notebooks. Skipping."
+                )
+        except ImportError:
+            logging.info(
+                "Failed to import `treescope` or `IPython`. Ignoring `treescope` registration"
+            )
+
+    log_handlers: list[logging.Handler] = []
+    if log_save_dir:
+        log_file = log_save_dir / "logging.log"
+        log_file.touch(exist_ok=True)
+        log_handlers.append(logging.FileHandler(log_file))
+
+    if rich:
+        try:
+            from rich.logging import RichHandler  # type: ignore
+
+            log_handlers.append(RichHandler(rich_tracebacks=rich_tracebacks))
+        except ImportError:
+            logging.info(
+                "Failed to import rich. Falling back to default Python logging."
+            )
+
+    logging.basicConfig(
+        level=log_level,
+        format="%(message)s",
+        datefmt="[%X]",
+        handlers=log_handlers,
+    )
+    logging.info(
+        "Logging initialized. "
+        f"Lovely: {lovely}, Treescope: {treescope}, Rich: {rich}, "
+        f"Log level: {log_level}, Log save dir: {log_save_dir}"
+    )
+
+
+init_python_logging = setup_logging

nshutils-0.20.0/src/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-0.20.0/src/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-0.20.0/src/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-0.20.0/src/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-0.20.0/src/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-0.20.0/src/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-0.20.0/src/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-0.20.0/src/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-0.19.0 → nshutils-0.20.0}/src/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-0.20.0/src/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.0/src/nshutils/logging.py

@@ -1,125 +0,0 @@
-from __future__ import annotations
-
-import logging
-from pathlib import Path
-
-
-def init_python_logging(
-    *,
-    lovely_tensors: bool = False,
-    lovely_numpy: bool = False,
-    treescope: bool = False,
-    treescope_autovisualize_arrays: bool = False,
-    rich: bool = False,
-    rich_tracebacks: bool = False,
-    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
-
-            _lovely_numpy.set_config(repr=_lovely_numpy.lovely)
-        except ImportError:
-            logging.info(
-                "Failed to import `lovely_numpy`. Ignoring pretty numpy array formatting"
-            )
-
-    if treescope:
-        try:
-            # Check if we're in a Jupyter environment
-            from IPython import get_ipython
-
-            if get_ipython() is not None:
-                import treescope as _treescope  # type: ignore
-
-                _treescope.basic_interactive_setup(
-                    autovisualize_arrays=treescope_autovisualize_arrays
-                )
-            else:
-                logging.info(
-                    "Treescope setup is only supported in Jupyter notebooks. Skipping."
-                )
-        except ImportError:
-            logging.info(
-                "Failed to import `treescope` or `IPython`. Ignoring `treescope` registration"
-            )
-
-    log_handlers: list[logging.Handler] = []
-    if log_save_dir:
-        log_file = log_save_dir / "logging.log"
-        log_file.touch(exist_ok=True)
-        log_handlers.append(logging.FileHandler(log_file))
-
-    if rich:
-        try:
-            from rich.logging import RichHandler  # type: ignore
-
-            log_handlers.append(RichHandler(rich_tracebacks=rich_tracebacks))
-        except ImportError:
-            logging.info(
-                "Failed to import rich. Falling back to default Python logging."
-            )
-
-    logging.basicConfig(
-        level=log_level,
-        format="%(message)s",
-        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,
-    )
-
-
-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,
-    )