code: def tree_pprint(
      pytree: PrettyPrintAble,
      *,
      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,
  ) -> None:
      """Pretty-prints a PyTree as a string, whilst abbreviating JAX arrays.
  
      All JAX arrays in the PyTree are condensed down to a short string representation
      of their dtype and shape.
  
      !!! example
  
          A 32-bit floating-point JAX array of shape `(3, 4)` is printed as `f32[3,4]`.
  
      **Arguments:**
  
      - `pytree`: The PyTree to pretty-print.
      - `width`: The desired maximum number of characters per line of output. If a
          structure cannot be formatted within this constraint then a best effort will
          be made.
      - `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__}(...)"`.
  
      **Returns:**
  
      Nothing. (The result is printed to stdout instead.)
      """
      print(
          tree_pformat(
              pytree,
              width=width,
              indent=indent,
              short_arrays=short_arrays,
              struct_as_array=struct_as_array,
              follow_wrapped=follow_wrapped,
              truncate_leaf=truncate_leaf,
          )
      )
  
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 16 for calling a global function with example usage: tree_pprint(pytree) without return types