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: Takes a smaller `State` object (typically produces via `.substate`), and
  updates states by using all of its values.
  
  **Arguments:**
  
  - `substate`: a `State` object whose keys are a subset of the keys of `self`.
  
  **Returns:**
  
  A new [`equinox.nn.State`][] object, containing all of the updated values.
  
  As a safety guard against accidentally writing `state.set(item, value)` without
  assigning it to a new value, then the old object (`self`) will become invalid.
docstr_tok: ["Takes","a","smaller","`","State","`","object","(","typically","produces","via","`",".substate","`",")",",","and","updates","states","by","using","all","of","its","values",".","*","*","Arguments",":","*","*","-","`","substate","`",":","a","`","State","`","object","whose","keys","are","a","subset","of","the","keys","of","`","self","`",".","*","*","Returns",":","*","*","A","new","[","`","equinox.nn.State","`","]","[","]","object",",","containing","all","of","the","updated","values",".","As","a","safety","guard","against","accidentally","writing","`","state.set","(","item",",","value",")","`","without","assigning","it","to","a","new","value",",","then","the","old","object","(","`","self","`",")","will","become","invalid","."]
code: def update(self, substate: "State") -> "State":
      """Takes a smaller `State` object (typically produces via `.substate`), and
      updates states by using all of its values.
  
      **Arguments:**
  
      - `substate`: a `State` object whose keys are a subset of the keys of `self`.
  
      **Returns:**
  
      A new [`equinox.nn.State`][] object, containing all of the updated values.
  
      As a safety guard against accidentally writing `state.set(item, value)` without
      assigning it to a new value, then the old object (`self`) will become invalid.
      """
      if self._state is _sentinel:
          raise ValueError(_state_error)
      if type(substate) is not State:
          raise ValueError("Can only use `eqx.nn.State`s in `update`.")
      state = self._state.copy()  # pyright: ignore
      for key, value in substate._state.items():  # pyright: ignore
          if key not in state:
              raise ValueError(
                  "Cannot call `state1.update(state2)` unless `state2` is a substate "
                  "of `state1`."
              )
          state[key] = value
      new_self = object.__new__(State)
      new_self._state = state
      self._state = _sentinel
      return new_self
  
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 5 for calling the State obj's underlying member method code with example usage: obj.update(substate) and returns: new_self