class_docstr: This wraps together (a) a unique dictionary key used for looking up a stateful
  value, and (b) how that stateful value should be initialised.
  
  !!! Example
  
      ```python
      class MyStatefulLayer(eqx.Module):
          index: eqx.nn.StateIndex
  
          def __init__(self):
              init_state = jnp.array(0)
              self.index = eqx.nn.StateIndex(init_state)
  
          def __call__(self, x: Array, state: eqx.nn.State) -> tuple[Array, eqx.nn.State]:
              current_state = state.get(self.index)
              new_x = x + current_state
              new_state = state.set(self.index, current_state + 1)
              return new_x, new_state
      ```
  
  See also e.g. the source code of built-in stateful layers like
  [`equinox.nn.BatchNorm`][] for further reference.
docstr: **Arguments:**
  
  - `init`: The initial value for the state.
docstr_tok: ["*","*","Arguments",":","*","*","-","`","init","`",":","The","initial","value","for","the","state","."]
code: def __init__(self, init: _Value):
      """**Arguments:**
  
      - `init`: The initial value for the state.
      """
      if isinstance(init, types.FunctionType):
          # Technically a function is valid here, since we could allow any pytree.
          # In practice that's weird / kind of useless, so better to explicitly raise
          # the deprecation error.
          raise ValueError(
              "As of Equinox v0.11.0, `eqx.nn.StateIndex` now accepts the value "
              "of the initial state directly. (Not a function that creates the "
              "initial state.)"
          )
      self.marker = object()
      self.init = init
  
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 StateIndex object from inherited base classes: Module, Generic with example: obj = StateIndex(init)