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: Creates a smaller `State` object, that tracks only the states of some smaller
  part of the overall model.
  
  **Arguments:**
  
  - `pytree`: any PyTree. It will be iterated over to check for
      [`equinox.nn.StateIndex`]s.
  
  **Returns:**
  
  A new [`equinox.nn.State`][] object, which tracks only some of the overall
  states.
docstr_tok: ["Creates","a","smaller","`","State","`","object",",","that","tracks","only","the","states","of","some","smaller","part","of","the","overall","model",".","*","*","Arguments",":","*","*","-","`","pytree","`",":","any","PyTree",".","It","will","be","iterated","over","to","check","for","[","`","equinox.nn.StateIndex","`","]","s.","*","*","Returns",":","*","*","A","new","[","`","equinox.nn.State","`","]","[","]","object",",","which","tracks","only","some","of","the","overall","states","."]
code: def substate(self, pytree: PyTree) -> "State":
      """Creates a smaller `State` object, that tracks only the states of some smaller
      part of the overall model.
  
      **Arguments:**
  
      - `pytree`: any PyTree. It will be iterated over to check for
          [`equinox.nn.StateIndex`]s.
  
      **Returns:**
  
      A new [`equinox.nn.State`][] object, which tracks only some of the overall
      states.
      """
      if self._state is _sentinel:
          raise ValueError(_state_error)
      leaves = jtu.tree_leaves(pytree, is_leaf=_is_index)
      markers = [x.marker for x in leaves if _is_index(x)]
      substate = {k: self._state[k] for k in markers}  # pyright: ignore
      subself = object.__new__(State)
      subself._state = substate
      return subself
  
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 4 for calling the State obj's underlying member method code with example usage: obj.substate(pytree) and returns: subself