code: def delete_init_state(model: PyTree) -> PyTree:
      """For memory efficiency, this deletes the initial state stored within a model.
  
      Every stateful layer keeps a copy of the initial value of its state. This is then
      collected by [`equinox.nn.State`][], when it is called on the model. However, this
      means that the model must keep a copy of the initial state around, in case
      `eqx.nn.State` is called on it again. This extra copy consumes extra memory.
  
      But in practice, it is quite common to only need to initialise the state once. In
      this case, we can use this function to delete this extra copy, and in doing so save
      some memory.
  
      !!! Example
  
          Here is the typical pattern in which this is used:
          ```python
          model_and_state = eqx.nn.BatchNorm(...)
          state = eqx.nn.State(model_and_state)
          model = eqx.nn.delete_init_state(model)
          del model_and_state  # ensure this goes out of scope and is garbage collected
          ```
          Indeed the above is precisely what [`equinox.nn.make_with_state`][] does.
  
      **Arguments:**
  
      - `model`: any PyTree.
  
      **Returns:**
  
      A copy of `model`, with all the initial states stripped out. (As in the exampels
      above, you should then dispose of the original `model` object.)
      """
      return jtu.tree_map(_delete_init_state, model, is_leaf=_is_index)
  
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 a global function with example usage: delete_init_state(model) and returns: jtu