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: Sets a new value for an [`equinox.nn.StateIndex`][], **and returns the
  updated state**.
  
  **Arguments:**
  
  - `item`: an [`equinox.nn.StateIndex`][].
  - `value`: the new value associated with that index.
  
  **Returns:**
  
  A new [`equinox.nn.State`][] object, with the update.
  
  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: ["Sets","a","new","value","for","an","[","`","equinox.nn.StateIndex","`","]","[","]",",","*","*","and","returns","the","updated","state","*","*",".","*","*","Arguments",":","*","*","-","`","item","`",":","an","[","`","equinox.nn.StateIndex","`","]","[","]",".","-","`","value","`",":","the","new","value","associated","with","that","index",".","*","*","Returns",":","*","*","A","new","[","`","equinox.nn.State","`","]","[","]","object",",","with","the","update",".","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 set(self, item: StateIndex[_Value], value: _Value) -> "State":
      """Sets a new value for an [`equinox.nn.StateIndex`][], **and returns the
      updated state**.
  
      **Arguments:**
  
      - `item`: an [`equinox.nn.StateIndex`][].
      - `value`: the new value associated with that index.
  
      **Returns:**
  
      A new [`equinox.nn.State`][] object, with the update.
  
      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(item) is not StateIndex:
          raise ValueError(
              "Can only use `eqx.nn.StateIndex`s as state keys."
          )
      old_value = self._state[item.marker]  # pyright: ignore
      value = jtu.tree_map(jnp.asarray, value)
      old_struct = jax.eval_shape(lambda: old_value)
      new_struct = jax.eval_shape(lambda: value)
      if tree_equal(old_struct, new_struct) is not True:
          old_repr = tree_pformat(old_struct, struct_as_array=True)
          new_repr = tree_pformat(new_struct, struct_as_array=True)
          raise ValueError(
              "Old and new values have different structures/shapes/dtypes. The old "
              f"value is {old_repr} and the new value is {new_repr}."
          )
      state = self._state.copy()  # pyright: ignore
      state[item.marker] = 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 3 for calling the State obj's underlying member method code with example usage: obj.set(item, value) and returns: new_self