code: def lift(
      mod: nn.Module,
      *,
      name: str,
  ) -> Callable[..., Any]:
      """Lifts a flax nn.Module into a Haiku transformed function.
  
      For a Flax Module (e.g. ``mod = nn.Dense(10)``), ``mod = lift(mod)`` allows
      you to run the call method of the module as if the module was a regular Haiku
      module.
  
      Parameters and state from the Flax module are registered with Haiku and become
      part of the params/state dictionaries (as returned from ``init``/``apply``).
  
      >>> def f(x):
      ...   # Create and "lift" a Flax module.
      ...   mod = hk.experimental.flax.lift(nn.Dense(300), name='dense')
      ...   x = mod(x)                  # Any params/state will be registered
      ...                               # with Haiku when applying the module.
      ...   x = jax.nn.relu(x)
      ...   x = hk.nets.MLP([100, 10])  # You can of course mix Haiku modules in.
      ...   return x
      >>> f = hk.transform(f)
      >>> rng = jax.random.PRNGKey(42)
      >>> x = jnp.ones([1, 1])
      >>> params = f.init(rng, x)
      >>> out = f.apply(params, None, x)
  
      Args:
        mod: Any Flax ``nn.Module`` instance.
        name: Name scope to prefix entries in the outer params/state dict.
  
      Returns:
        A function that when applied calls the call method of the given Flax module
        and returns its output. As a side effect of calling the module any module
        parameters and state variables are registered with Haiku.
      """
      mod = _flax_transform_with_state(mod)
      init_fn, updater = hk.lift_with_state(
          mod.init, name=name, allow_reuse=True
      )
  
      def wrapped(*args, **kwargs):
          init_kwargs = dict(kwargs)
          if hk.running_init():
              if "rngs" not in init_kwargs:
                  init_rng = hk.maybe_next_rng_key()
              elif isinstance(init_kwargs["rngs"], Mapping):
                  init_rng = dict(init_kwargs.pop("rngs"))
                  if "params" not in init_rng:
                      rng = hk.maybe_next_rng_key()
                      if rng is not None:
                          init_rng["params"] = rng
              else:
                  raise flax.errors.InvalidRngError(
                      "rngs should be a dictionary mapping strings to `jax.PRNGKey`."
                  )
          else:
              init_rng = None
          params, state = init_fn(init_rng, *args, **init_kwargs)
          out, state = mod.apply(params, state, None, *args, **kwargs)
          updater._used = False  # pylint: disable=protected-access
          updater.update(state)
          return out
  
      return wrapped
  
file_path: dm-haiku/haiku/_src/flax/transform_flax.py
imports: from collections.abc import Mapping
  from typing import Any, Callable
  import flax.errors
  import flax.linen
  from haiku._src import base
  from haiku._src import filtering
  from haiku._src import lift
  from haiku._src import transform
  from haiku._src import typing
  from haiku._src.flax import utils
label_desc: Use image node_id 5 for calling a global function with example usage: lift(mod) and returns: wrapped, out