code: def _flax_transform_with_state(
      mod: nn.Module,
  ) -> hk.TransformedWithState:
      """Transforms a Flax ``nn.Module`` into a Haiku transformed function.
  
      Example usage:
  
      >>> mod = nn.Dense(10)
      >>> f = _flax_transform_with_state(mod)
      >>> rng = jax.random.PRNGKey(42)
      >>> x = jnp.ones([1, 1])
      >>> params, state = f.init(rng, x)
      >>> out, state = f.apply(params, state, rng, x)
  
      Args:
        mod: Any Flax ``nn.Module`` instance.
  
      Returns:
        A :class:`~haiku.TransformedWithState` instance (equivalent to the result of
        calling :func:`transform_with_state`).
  
      See also:
        lift: Use a Flax module as part of an outer :func:`~haiku.transform` or
          :func:`~haiku.transform_with_state`.
      """
  
      def init_fn(rng, *args, **kwargs):
          assert "rngs" not in kwargs  # Handled in `lift`.
          variables = dict(mod.init(rng, *args, **kwargs))
          params = utils.flatten_flax_to_haiku(
              variables.pop("params", {})
          )
          state = _to_haiku_state(variables)
          return params, state
  
      def apply_fn(params, state, rng, *args, **kwargs):
          if rng is not None:
              raise ValueError(
                  "RNGs passed in apply must be passed in the rngs keyword argument"
              )
  
          variables = _from_haiku_state(state)
          mutable = set(variables)
          variables["params"] = _from_haiku_params(params)
          out, variables = mod.apply(
              variables, *args, **kwargs, mutable=mutable
          )
          state = _to_haiku_state(variables)
          return out, state
  
      return hk.TransformedWithState(init_fn, apply_fn)
  
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 4 for calling a global function with example usage: _flax_transform_with_state(mod) and returns: hk, params,  state, out,  state