code: 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))
  
imports: import dataclasses
  import functools
  import sys
  import types
  from collections.abc import Callable, Sequence
  from typing import Any, Optional, Union
  import jax
  import jax._src.pretty_printer
  import jax.tree_util
  import numpy
  from jaxtyping import PyTree
  from ._doc_utils import WithRepr
label_desc: Use image node_id 13 for calling a global function with example usage: tree_pp(obj) and returns: pp, _pformat_dataclass, pp, _pformat_list, _pformat_dict, _pformat_namedtuple, _pformat_tuple, _pformat_array, tree_pp, tree_pp, _pformat_dataclass, tree_pp, _pformat_function, pp