code: def transform_with_state(f) -> TransformedWithState:
      """Transforms a function using Haiku modules into a pair of pure functions.
  
      See :func:`transform` for general details on Haiku transformations.
  
      For a function ``out = f(*a, **k)`` this function returns a pair of two pure
      functions that call ``f(*a, **k)`` explicitly collecting and injecting
      parameter values and state::
  
          params, state = init(rng, *a, **k)
          out, state = apply(params, state, rng, *a, **k)
  
      Note that the ``rng`` argument is typically not required for ``apply`` and
      passing ``None`` is accepted.
  
      This function is equivalent to :func:`transform`, however it allows you to
      maintain and update internal state (e.g. :class:`ExponentialMovingAverage` in
      :class:`BatchNorm`) via :func:`get_state` and :func:`set_state`:
  
      >>> def f():
      ...   counter = hk.get_state("counter", shape=[], dtype=jnp.int32,
      ...                          init=jnp.zeros)
      ...   hk.set_state("counter", counter + 1)
      ...   return counter
      >>> f = hk.transform_with_state(f)
  
      >>> params, state = f.init(None)
      >>> for _ in range(10):
      ...   counter, state = f.apply(params, state, None)
      >>> print(counter)
      9
  
      Args:
        f: A function closing over :class:`Module` instances.
  
      Returns:
        A :class:`TransformedWithState` tuple with ``init`` and ``apply`` pure
        functions.
      """
      analytics.log_once("transform_with_state")
      check_not_jax_transformed(f)
  
      unexpected_tracer_hint = (
          "An UnexpectedTracerError was raised while inside a Haiku transformed "
          "function (see error above).\n"
          "Hint: are you using a JAX transform or JAX control-flow function "
          "(jax.vmap/jax.lax.scan/...) inside a Haiku transform? You might want to use "
          "the Haiku version of the transform instead (hk.vmap/hk.scan/...).\n"
          "See https://dm-haiku.readthedocs.io/en/latest/notebooks/transforms.html "
          "on why you can't use JAX transforms inside a Haiku module."
      )
  
      f_sig = inspect.signature(f)
  
      def init_fn(
          rng: Optional[Union[PRNGKey, int]],
          *args,
          **kwargs,
      ) -> tuple[hk.MutableParams, hk.MutableState]:
          """Initializes your function collecting parameters and state."""
          rng = to_prng_sequence(rng, err_msg=INIT_RNG_ERROR)
          with base.new_context(rng=rng) as ctx:
              try:
                  f(*args, **kwargs)
              except jax.errors.UnexpectedTracerError as e:
                  raise jax.errors.UnexpectedTracerError(
                      unexpected_tracer_hint
                  ) from e
          return ctx.collect_params(), ctx.collect_initial_state()
  
      init_fn.__signature__ = inspect.Signature(
          parameters=[
              inspect.Parameter(
                  "rng",
                  inspect.Parameter.POSITIONAL_OR_KEYWORD,
                  annotation=Optional[Union[PRNGKey, int]],
              ),
          ]
          + list(f_sig.parameters.values()),
          return_annotation=tuple[hk.Params, hk.State],
          __validate_parameters__=False,
      )
  
      def apply_fn(
          params: Optional[hk.Params],
          state: Optional[hk.State],
          rng: Optional[Union[PRNGKey, int]],
          *args,
          **kwargs,
      ) -> tuple[Any, hk.MutableState]:
          """Applies your function injecting parameters and state."""
          uses_state = state is not None
          params = check_mapping("params", params)
          state = check_mapping("state", state)
          rng = to_prng_sequence(
              rng,
              err_msg=(
                  APPLY_RNG_STATE_ERROR
                  if uses_state
                  else APPLY_RNG_ERROR
              ),
          )
          with base.new_context(
              params=params, state=state, rng=rng
          ) as ctx:
              try:
                  out = f(*args, **kwargs)
              except jax.errors.UnexpectedTracerError as e:
                  raise jax.errors.UnexpectedTracerError(
                      unexpected_tracer_hint
                  ) from e
          return out, ctx.collect_state()
  
      apply_fn.__signature__ = sig_add_state(
          inspect.Signature(
              parameters=[
                  inspect.Parameter(
                      "params",
                      inspect.Parameter.POSITIONAL_OR_KEYWORD,
                      annotation=Optional[hk.Params],
                  ),
                  inspect.Parameter(
                      "state",
                      inspect.Parameter.POSITIONAL_OR_KEYWORD,
                      annotation=Optional[hk.State],
                  ),
                  inspect.Parameter(
                      "rng",
                      inspect.Parameter.POSITIONAL_OR_KEYWORD,
                      annotation=Optional[Union[PRNGKey, int]],
                  ),
              ]
              + list(f_sig.parameters.values()),
              return_annotation=f_sig.return_annotation,
              __validate_parameters__=False,
          )
      )
  
      tie_in_original_fn(f, init_fn, apply_fn)
  
      return TransformedWithState(init_fn, apply_fn)
  
imports: from collections.abc import Mapping
  import inspect
  from typing import Any, Callable, NamedTuple, Optional, TypeVar, Union
  from haiku._src import analytics
  from haiku._src import base
  from haiku._src import data_structures
  from haiku._src import typing
  import jax
label_desc: Use image node_id 9 for calling a global function with example usage: transform_with_state(f) and returns: TransformedWithState, ctx,  ctx, out,  ctx