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`][].
docstr: **Arguments:**
  
  - `model`: any PyTree. All stateful layers (e.g. [`equinox.nn.BatchNorm`][])
      will have their initial state stored inside the `State` object.
docstr_tok: ["*","*","Arguments",":","*","*","-","`","model","`",":","any","PyTree",".","All","stateful","layers","(","e.g",".","[","`","equinox.nn.BatchNorm","`","]","[","]",")","will","have","their","initial","state","stored","inside","the","`","State","`","object","."]
code: def __init__(self, model: PyTree):
      """**Arguments:**
  
      - `model`: any PyTree. All stateful layers (e.g. [`equinox.nn.BatchNorm`][])
          will have their initial state stored inside the `State` object.
      """
      # Note that de/serialisation depends on the ordered-ness of this dictionary,
      # between serialisation and deserialisation.
      state = {}
      leaves = jtu.tree_leaves(model, is_leaf=_is_index)
      for leaf in leaves:
          if _is_index(leaf):
              if leaf.init is _sentinel:
                  raise ValueError(
                      "Cannot call `eqx.nn.State(eqx.nn.delete_init_state(model))`. "
                      "You should call `eqx.nn.State(model)`, using the original "
                      "model."
                  )
              state[leaf.marker] = jtu.tree_map(jnp.asarray, leaf.init)
      self._state = state
  
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 1 to create a new State object  with example: obj = State(model)