nshutils 0.22.1__tar.gz → 0.22.3__tar.gz

{nshutils-0.22.1 → nshutils-0.22.3}/PKG-INFO

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

{nshutils-0.22.1 → nshutils-0.22.3}/pyproject.toml

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

{nshutils-0.22.1 → nshutils-0.22.3}/src/nshutils/lovely/_monkey_patch_all.py

@@ -41,21 +41,20 @@ def monkey_patch(libraries: list[Library] | Literal["auto"] = "auto"):
 
     with contextlib.ExitStack() as stack:
         for library in libraries:
-            match library:
-                case "torch":
-                    from .torch_ import torch_monkey_patch
+            if library == "torch":
+                from .torch_ import torch_monkey_patch
 
-                    stack.enter_context(torch_monkey_patch())
-                case "jax":
-                    from .jax_ import jax_monkey_patch
+                stack.enter_context(torch_monkey_patch())
+            elif library == "jax":
+                from .jax_ import jax_monkey_patch
 
-                    stack.enter_context(jax_monkey_patch())
-                case "numpy":
-                    from .numpy_ import numpy_monkey_patch
+                stack.enter_context(jax_monkey_patch())
+            elif library == "numpy":
+                from .numpy_ import numpy_monkey_patch
 
-                    stack.enter_context(numpy_monkey_patch())
-                case _:
-                    assert_never(library)
+                stack.enter_context(numpy_monkey_patch())
+            else:
+                assert_never(library)
 
         log.info(
             f"Monkey patched libraries: {', '.join(libraries)}. "

nshutils-0.22.3/src/nshutils/lovely/numpy_.py

@@ -0,0 +1,109 @@
+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 _np_ge_2():
+    import importlib.metadata
+
+    from packaging.version import Version
+
+    try:
+        numpy_version = importlib.metadata.version("numpy")
+        return Version(numpy_version) >= Version("2.0")
+    except importlib.metadata.PackageNotFoundError:
+        return False
+
+
+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),
+    }
+
+
+# If numpy 2.0, use the new API override_repr.
+if _np_ge_2():
+
+    @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()=}"
+            )
+
+else:
+
+    @monkey_patch_contextmanager(dependencies=["numpy"])
+    def numpy_monkey_patch():
+        try:
+            np.set_string_function(numpy_repr, True)  # pyright: ignore[reportAttributeAccessIssue]
+            np.set_string_function(numpy_repr, False)  # pyright: ignore[reportAttributeAccessIssue]
+
+            logging.info(
+                f"Numpy monkey patching: using {numpy_repr.__name__} for numpy arrays. "
+                f"{np.get_printoptions()=}"
+            )
+            yield
+        finally:
+            np.set_string_function(None, True)  # pyright: ignore[reportAttributeAccessIssue]
+            np.set_string_function(None, False)  # pyright: ignore[reportAttributeAccessIssue]
+
+            logging.info(
+                f"Numpy unmonkey patching: using {numpy_repr.__name__} for numpy arrays. "
+                f"{np.get_printoptions()=}"
+            )

nshutils-0.22.1/src/nshutils/lovely/numpy_.py

@@ -1,72 +0,0 @@
-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()=}"
-        )