Now using wadler_lindig pprint library.

Author: patrick-kidger

.pre-commit-config.yaml

@@ -20,4 +20,5 @@ repos:
             optax,
             pytest,
             typing_extensions,
+            wadler_lindig,
           ]

README.md

@@ -19,7 +19,7 @@ _Coming from [Flax](https://github.com/google/flax) or [Haiku](https://github.co
 pip install equinox
 ```
 
-Requires Python 3.9+ and JAX 0.4.13+.
+Requires Python 3.10+.
 
 ## Documentation
 

docs/index.md

@@ -19,7 +19,7 @@ _Coming from [Flax](https://github.com/google/flax) or [Haiku](https://github.co
 pip install equinox
 ```
 
-Requires Python 3.9+ and JAX 0.4.13+.
+Requires Python 3.10+.
 
 ## Quick example
 

equinox/_pretty_print.py

@@ -1,183 +1,10 @@
 import dataclasses
-import functools as ft
-import sys
-import types
-from collections.abc import Callable, Sequence
-from typing import Any, Optional, Union
+from collections.abc import Callable
+from typing import Any
 
 import jax
-import jax._src.pretty_printer as pp
 import jax.tree_util as jtu
-import numpy as np
-from jaxtyping import PyTree
-
-from ._doc_utils import WithRepr
-
-
-Dataclass = Any
-NamedTuple = Any  # workaround typeguard bug
-PrettyPrintAble = PyTree
-
-# Re-export
-text = pp.text
-
-
-_comma_sep = pp.concat([pp.text(","), pp.brk()])
-
-
-def bracketed(
-    name: Optional[pp.Doc],
-    indent: int,
-    objs: Sequence[pp.Doc],
-    lbracket: str,
-    rbracket: str,
-) -> pp.Doc:
-    nested = pp.concat(
-        [
-            pp.nest(indent, pp.concat([pp.brk(""), pp.join(_comma_sep, objs)])),
-            pp.brk(""),
-        ]
-    )
-    concated = []
-    if name is not None:
-        concated.append(name)
-    concated.extend([pp.text(lbracket), nested, pp.text(rbracket)])
-    return pp.group(pp.concat(concated))
-
-
-def named_objs(pairs, **kwargs):
-    return [
-        pp.concat([pp.text(key + "="), tree_pp(value, **kwargs)])
-        for key, value in pairs
-    ]
-
-
-def _pformat_list(obj: list, **kwargs) -> pp.Doc:
-    return bracketed(
-        name=None,
-        indent=kwargs["indent"],
-        objs=[tree_pp(x, **kwargs) for x in obj],
-        lbracket="[",
-        rbracket="]",
-    )
-
-
-def _pformat_tuple(obj: tuple, **kwargs) -> pp.Doc:
-    if len(obj) == 1:
-        objs = [pp.concat([tree_pp(obj[0], **kwargs), pp.text(",")])]
-    else:
-        objs = [tree_pp(x, **kwargs) for x in obj]
-    return bracketed(
-        name=None, indent=kwargs["indent"], objs=objs, lbracket="(", rbracket=")"
-    )
-
-
-def _pformat_namedtuple(obj: NamedTuple, **kwargs) -> pp.Doc:
-    objs = named_objs([(name, getattr(obj, name)) for name in obj._fields], **kwargs)
-    return bracketed(
-        name=pp.text(obj.__class__.__name__),
-        indent=kwargs["indent"],
-        objs=objs,
-        lbracket="(",
-        rbracket=")",
-    )
-
-
-def _dict_entry(key: Any, value: Any, **kwargs) -> pp.Doc:
-    return pp.concat(
-        [tree_pp(key, **kwargs), pp.text(":"), pp.brk(), tree_pp(value, **kwargs)]
-    )
-
-
-def _pformat_dict(obj: dict, **kwargs) -> pp.Doc:
-    objs = [_dict_entry(key, value, **kwargs) for key, value in obj.items()]
-    return bracketed(
-        name=None,
-        indent=kwargs["indent"],
-        objs=objs,
-        lbracket="{",
-        rbracket="}",
-    )
-
-
-def pformat_short_array_text(shape: tuple[int, ...], dtype: str) -> str:
-    short_dtype = (
-        dtype.replace("float", "f")
-        .replace("uint", "u")
-        .replace("int", "i")
-        .replace("complex", "c")
-    )
-    short_shape = ",".join(map(str, shape))
-    return f"{short_dtype}[{short_shape}]"
-
-
-def _pformat_short_array(
-    shape: tuple[int, ...], dtype: str, kind: Optional[str]
-) -> pp.Doc:
-    out = pformat_short_array_text(shape, dtype)
-    if kind is not None:
-        out = out + f"({kind})"
-    return pp.text(out)
-
-
-def _pformat_array(
-    obj: Union[
-        jax.Array,
-        jax.ShapeDtypeStruct,
-        np.ndarray,
-        "torch.Tensor",  # pyright: ignore  # noqa: F821
-    ],
-    **kwargs,
-) -> pp.Doc:
-    short_arrays = kwargs["short_arrays"]
-    if short_arrays:
-        # Support torch here for the sake of jaxtyping's pretty-printed error messages.
-        if "torch" in sys.modules and isinstance(obj, sys.modules["torch"].Tensor):
-            dtype = repr(obj.dtype).split(".")[1]
-            kind = "torch"
-        else:
-            dtype = obj.dtype.name
-            if isinstance(obj, (jax.Array, jax.ShapeDtypeStruct)):
-                # Added in JAX 0.4.32 to `ShapeDtypeStruct`
-                if getattr(obj, "weak_type", False):
-                    dtype = f"weak_{dtype}"
-                kind = None
-            elif isinstance(obj, np.ndarray):
-                kind = "numpy"
-            else:
-                kind = "unknown"
-        return _pformat_short_array(obj.shape, dtype, kind)
-    else:
-        return pp.text(repr(obj))
-
-
-def _pformat_function(obj: types.FunctionType, **kwargs) -> pp.Doc:
-    if kwargs.get("wrapped", False):
-        fn = "wrapped function"
-    else:
-        fn = "function"
-    return pp.text(f"<{fn} {obj.__name__}>")
-
-
-def _pformat_dataclass(obj, **kwargs) -> pp.Doc:
-    # <uninitialised> can happen when typechecking an `eqx.Module`'s `__init__` method
-    # with beartype, and printing args using pytest. We haven't yet actually assigned
-    # values to the module so the repr fails.
-    objs = named_objs(
-        [
-            (field.name, getattr(obj, field.name, WithRepr(None, "<uninitialised>")))
-            for field in dataclasses.fields(obj)
-            if field.repr
-        ],
-        **kwargs,
-    )
-    return bracketed(
-        name=pp.text(obj.__class__.__name__),
-        indent=kwargs["indent"],
-        objs=objs,
-        lbracket="(",
-        rbracket=")",
-    )
+import wadler_lindig as wl
 
 
 @dataclasses.dataclass
@@ -192,63 +19,18 @@ class _Partial:
 _Partial.__module__ = jtu.Partial.__module__
 
 
-def tree_pp(obj: PrettyPrintAble, **kwargs) -> pp.Doc:
-    follow_wrapped = kwargs["follow_wrapped"]
-    truncate_leaf = kwargs["truncate_leaf"]
-    if truncate_leaf(obj):
-        return pp.text(f"{type(obj).__name__}(...)")
-    if hasattr(obj, "__tree_pp__"):
-        custom_pp = obj.__tree_pp__(**kwargs)
-        if custom_pp is not NotImplemented:
-            return pp.group(custom_pp)
-    if dataclasses.is_dataclass(obj) and not isinstance(obj, type):
-        return _pformat_dataclass(obj, **kwargs)
-    elif isinstance(obj, list):
-        return _pformat_list(obj, **kwargs)
-    elif isinstance(obj, dict):
-        return _pformat_dict(obj, **kwargs)
-    elif isinstance(obj, tuple):
-        if hasattr(obj, "_fields"):
-            return _pformat_namedtuple(obj, **kwargs)
-        else:
-            return _pformat_tuple(obj, **kwargs)
-    elif (
-        isinstance(obj, (np.ndarray, jax.Array))
-        or ("torch" in sys.modules and isinstance(obj, sys.modules["torch"].Tensor))
-        or kwargs.get("struct_as_array", False)
-        and isinstance(obj, jax.ShapeDtypeStruct)
-    ):
-        return _pformat_array(obj, **kwargs)
-    elif isinstance(obj, (jax.custom_jvp, jax.custom_vjp)):
-        return tree_pp(obj.__wrapped__, **kwargs)
-    elif hasattr(obj, "__wrapped__") and follow_wrapped:
-        kwargs["wrapped"] = True
-        return tree_pp(obj.__wrapped__, **kwargs)  # pyright: ignore
-    elif isinstance(obj, jtu.Partial) and follow_wrapped:
-        obj = _Partial(obj.func, obj.args, obj.keywords)
-        return _pformat_dataclass(obj, **kwargs)
-    elif isinstance(obj, ft.partial) and follow_wrapped:
-        kwargs["wrapped"] = True
-        return tree_pp(obj.func, **kwargs)
-    elif isinstance(obj, types.FunctionType):
-        return _pformat_function(obj, **kwargs)
-    else:  # int, str, float, complex, bool, etc.
-        return pp.text(repr(obj))
-
-
 def _false(_):
     return False
 
 
 def tree_pformat(
-    pytree: PrettyPrintAble,
+    pytree: Any,
     *,
     width: int = 80,
     indent: int = 2,
     short_arrays: bool = True,
     struct_as_array: bool = False,
-    follow_wrapped: bool = True,
-    truncate_leaf: Callable[[PrettyPrintAble], bool] = _false,
+    truncate_leaf: Callable[[Any], bool] = _false,
 ) -> str:
     """Pretty-formats a PyTree as a string, whilst abbreviating JAX arrays.
 
@@ -257,25 +39,40 @@ def tree_pformat(
     As [`equinox.tree_pprint`][], but returns the string instead of printing it.
     """
 
-    return tree_pp(
-        pytree,
-        indent=indent,
-        short_arrays=short_arrays,
-        struct_as_array=struct_as_array,
-        follow_wrapped=follow_wrapped,
-        truncate_leaf=truncate_leaf,
-    ).format(width=width)
+    def custom(obj):
+        if truncate_leaf(obj):
+            return wl.TextDoc(f"{type(obj).__name__}(...)")
+
+        if short_arrays:
+            if isinstance(obj, jax.Array) or (
+                struct_as_array and isinstance(obj, jax.ShapeDtypeStruct)
+            ):
+                dtype = obj.dtype.name
+                # Added in JAX 0.4.32 to `ShapeDtypeStruct`
+                if getattr(obj, "weak_type", False):
+                    dtype = f"weak_{dtype}"
+                return wl.array_summary(obj.shape, dtype, kind=None)
+
+        if isinstance(obj, (jax.custom_jvp, jax.custom_vjp)):
+            return wl.pdoc(obj.__wrapped__)
+
+        if isinstance(obj, jtu.Partial):
+            obj = _Partial(obj.func, obj.args, obj.keywords)
+            return wl.pdoc(obj)
+
+    return wl.pformat(
+        pytree, width=width, indent=indent, short_arrays=short_arrays, custom=custom
+    )
 
 
 def tree_pprint(
-    pytree: PrettyPrintAble,
+    pytree: Any,
     *,
     width: int = 80,
     indent: int = 2,
     short_arrays: bool = True,
     struct_as_array: bool = False,
-    follow_wrapped: bool = True,
-    truncate_leaf: Callable[[PrettyPrintAble], bool] = _false,
+    truncate_leaf: Callable[[Any], bool] = _false,
 ) -> None:
     """Pretty-prints a PyTree as a string, whilst abbreviating JAX arrays.
 
@@ -295,7 +92,6 @@ def tree_pprint(
     - `indent`: The amount of indentation each nesting level.
     - `short_arrays`: Toggles the abbreviation of JAX arrays.
     - `struct_as_array`: Whether to treat `jax.ShapeDtypeStruct`s as arrays.
-    - `follow_wrapped`: Whether to unwrap `functools.partial` and `functools.wraps`.
     - `truncate_leaf`: A function `Any -> bool`. Applied to all nodes in the PyTree;
         all truthy nodes will be truncated to just `f"{type(node).__name__}(...)"`.
 
@@ -310,7 +106,6 @@ def tree_pprint(
             indent=indent,
             short_arrays=short_arrays,
             struct_as_array=struct_as_array,
-            follow_wrapped=follow_wrapped,
             truncate_leaf=truncate_leaf,
         )
     )

equinox/debug/_dce.py

@@ -4,11 +4,12 @@
 import jax.core
 import jax.numpy as jnp
 import jax.tree_util as jtu
+import wadler_lindig as wl
 from jaxtyping import PyTree
 
 from .._doc_utils import WithRepr
 from .._filters import combine, is_array, partition
-from .._pretty_print import pformat_short_array_text, tree_pprint
+from .._pretty_print import tree_pprint
 
 
 _dce_store = {}
@@ -105,7 +106,7 @@ def inspect_dce(name: Hashable = None):
             except KeyError:
                 value = "<DCE'd>"
             else:
-                value = pformat_short_array_text(shape, dtype)
+                value = wl.array_summary(shape, dtype, kind=None).text
             new_leaves.append(WithRepr(None, value))
         tree = combine(jtu.tree_unflatten(treedef, new_leaves), static)
         print(f"Entry {i}:")

equinox/internal/__init__.py

@@ -21,7 +21,6 @@
 from .._eval_shape import cached_filter_eval_shape as cached_filter_eval_shape
 from .._misc import left_broadcast_to as left_broadcast_to
 from .._module import Static as Static
-from .._pretty_print import tree_pp as tree_pp
 from .._unvmap import (
     unvmap_all as unvmap_all,
     unvmap_all_p as unvmap_all_p,

equinox/internal/_str2jax.py

@@ -1,13 +1,15 @@
+import wadler_lindig as wl
+
 from .._module import Module
-from .._pretty_print import text as pp_text
 
 
 def str2jax(msg: str):
     """Creates a JAXable object whose `str(...)` is the specified string."""
 
     class M(Module):
-        def __tree_pp__(self, **kwargs):
-            return pp_text(msg)
+        def __pdoc__(self, **kwargs):
+            del kwargs
+            return wl.TextDoc(msg)
 
         def __repr__(self):
             return msg