code: def lift_with_state(
      init_fn: Callable[..., tuple[hk.Params, hk.State]],
      *,
      allow_reuse: bool = False,
      name: str = "lifted",
  ) -> tuple[
      Callable[..., tuple[hk.Params, hk.State]], LiftWithStateUpdater
  ]:
      r"""Registers params and state from an init function in an outer transform.
  
      See :func:`lift`\ for more context on when to use ``lift``.
  
      This function returns two objects. The first is a callable that runs your init
      function with slightly different behaviour based on if it's run during init
      vs. apply time. The second is an updater that can be used to pass updated
      state values that result from running your apply function. See later in the
      docs for a worked example.
  
      During init, the returned callable will run the given ``init_fn``, and include
      the resulting params/state in the outer transform's dictionaries. During
      ``apply``, the returned callable will instead pull the relevant params/state
      from the outer transform's dictionaries.
  
      Must be called inside :func:`transform_with_state`\ , and be passed the
      ``init`` member of a :class:`TransformedWithState`\ .
  
      By default, users must ensure that the given ``init`` does not accidentally
      catch modules from an outer :func:`transform_with_state` via functional
      closure. If this behavior is desirable, set ``allow_reuse`` to ``True``.
  
      Example:
  
        >>> def g(x):
        ...   return hk.nets.ResNet50(1)(x, True)
        >>> g = hk.transform_with_state(g)
        >>> params_and_state_fn, updater = (
        ...   hk.lift_with_state(g.init, name='f_lift'))
        >>> init_rng = hk.next_rng_key() if hk.running_init() else None
        >>> x = jnp.ones([1, 224, 224, 3])
        >>> params, state = params_and_state_fn(init_rng, x)
        >>> out, state = g.apply(params, state, None, x)
        >>> updater.update(state)
  
      Args:
        init_fn: The ``init`` function from an :class:`TransformedWithState`\ .
        allow_reuse: Allows lifted parameters and state to be reused from the
          outer :func:`transform_with_state`. This can be desirable when using
          ``lift_with_state`` within control flow (e.g. ``hk.scan``).
        name: A string name to prefix parameters with.
  
      Returns:
        A callable that during ``init`` injects parameter values into the outer
        context and during ``apply`` reuses parameters from the outer context. In
        both cases returns parameter values to be used with an ``apply`` function.
        The ``init`` function additionally returns an object used to update the
        outer context with new state after ``apply`` is called.
  
      See also:
        - :func:`lift`: Register parameters with an outer transform.
        - :func:`transparent_lift`: Register parameters with an outer transform
          without a namespace.
        - :func:`transparent_lift_with_state`: Register parameters and state with an
          outer transform without a namespace.
      """
      base.assert_context("lift_with_state")
      params_and_state_fn = _to_callable(
          LiftingModule(init_fn, allow_reuse=allow_reuse, name=name)
      )
      if base.current_module():
          name = f"{base.current_name()}/{name}"
      updater = LiftWithStateUpdater(name)
      return params_and_state_fn, updater
  
imports: from collections.abc import Mapping, MutableMapping
  import functools
  from typing import Any, Callable, Optional, TypeVar
  from haiku._src import base
  from haiku._src import data_structures
  from haiku._src import module
  from haiku._src import transform
  from haiku._src.typing import LiftingModuleType
label_desc: Use image node_id 8 for calling a global function with example usage: lift_with_state(init_fn) and returns: params_and_state_fn,  updater