class_docstr: Stores the state of a model. For example, the running statistics of all
  [`equinox.nn.BatchNorm`][] layers in the model.
  
  This is essentially a dictionary mapping from [`equinox.nn.StateIndex`][]s to
  PyTrees of arrays.
  
  This class should be initialised via [`equinox.nn.make_with_state`][].
code: def __tree_pp__(self, **kwargs):
      if self._state is _sentinel:
          return text("State(~old~)")
      else:
          objs = named_objs(
              [
                  (hex(id(key)), value)
                  for key, value in self._state.items()  # pyright: ignore
              ],
              **kwargs,
          )
          return bracketed(
              name=text("State"),
              indent=kwargs["indent"],
              objs=objs,
              lbracket="(",
              rbracket=")",
          )
  
imports: import types
  from collections.abc import Callable
  from typing import Any, Generic, TYPE_CHECKING, TypeVar, Union
  from typing_extensions import ParamSpec
  import jax
  import jax.numpy
  import jax.tree_util
  from jaxtyping import PyTree
  from .._module import field, Module
  from .._pretty_print import bracketed, named_objs, text, tree_pformat
  from .._tree import tree_at, tree_equal
label_desc: Use image node_id 7 for calling the State obj's underlying member method code with example usage: obj.__tree_pp__() and returns: text, bracketed