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_flatten(self):
      if self._state is _sentinel:
          raise ValueError(_state_error)
      keys = tuple(self._state.keys())  # pyright: ignore
      values = tuple(self._state[k] for k in keys)  # pyright: ignore
      return values, keys
  
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 8 for calling the State obj's underlying member method code with example usage: obj.tree_flatten() and returns: values,  keys