nshutils 0.19.1__tar.gz → 0.21.0__tar.gz

{nshutils-0.19.1 → nshutils-0.21.0}/PKG-INFO

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

@@ -1,15 +1,25 @@
 [project]
 name = "nshutils"
-version = "0.19.1"
+version = "0.21.0"
 description = ""
 authors = [{ name = "Nima Shoghi", email = "nimashoghi@gmail.com" }]
-requires-python = ">=3.10,<4.0"
+requires-python = ">=3.9,<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; python_version >= '3.10'"]
+extra = ["pysnooper", "rich[jupyter]", "treescope; python_version >= '3.10'"]
 
 [project.urls]
 homepage = "https://github.com/nimashoghi/nshutils"

nshutils-0.21.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.1/src/nshutils/__init__.py → nshutils-0.21.0/src/nshutils/__init__.pyi

@@ -1,29 +1,13 @@
-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
 from .typecheck import typecheck_this_module as typecheck_this_module
-
-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.1 → nshutils-0.21.0}/src/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-0.21.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.21.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.21.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.21.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.21.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.21.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)