code: def make_with_state(
      make_model: Callable[_P, _T]
  ) -> Callable[_P, tuple[_T, State]]:
      """This function is the most common API for working with stateful models. This
      initialises both the parameters and the state of a stateful model.
  
      `eqx.nn.make_with_state(Model)(*args, **kwargs)` simply calls
      `model_with_state = Model(*args, **kwargs)`, and then partitions the resulting
      PyTree into two pieces: the parameters, and the state.
  
      **Arguments:**
  
      - `make_model`: some callable returning a PyTree.
  
      **Returns:**
  
      A callable, which when evaluated returns a 2-tuple of `(model, state)`, where
      `model` is the result of `make_model(*args, **kwargs)` but with all of the initial
      states stripped out, and `state` is an [`equinox.nn.State`][] object encapsulating
      the initial states.
  
      !!! Example
  
          See [the stateful example](../../examples/stateful.ipynb) for a runnable
          example.
  
          ```python
          class Model(eqx.Module):
              def __init__(self, foo, bar):
                  ...
  
              ...
  
          model, state = eqx.nn.make_with_state(Model)(foo=3, bar=4)
          ```
      """
  
      # _P.{args, kwargs} not supported by beartype
      if TYPE_CHECKING:
  
          def make_with_state_impl(
              *args: _P.args, **kwargs: _P.kwargs
          ) -> tuple[_T, State]:
              ...
  
      else:
  
          def make_with_state_impl(*args, **kwargs) -> tuple[_T, State]:
              model = make_model(*args, **kwargs)
  
              # Replace all markers with `int`s. This is needed to ensure that two calls
              # to `make_with_state` produce compatible models and states.
              leaves, treedef = jtu.tree_flatten(
                  model, is_leaf=_is_index
              )
              counter = 0
              new_leaves = []
              for leaf in leaves:
                  if _is_index(leaf):
                      leaf = StateIndex(leaf.init)
                      object.__setattr__(leaf, "marker", counter)
                      counter += 1
                  new_leaves.append(leaf)
              model = jtu.tree_unflatten(treedef, new_leaves)
  
              state = State(model)
              model = delete_init_state(model)
              return model, state
  
      return make_with_state_impl
  
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 a global function with example usage: make_with_state(make_model) and returns: make_with_state_impl, model,  state