code: def transform(f, *, apply_rng=True) -> Transformed:
      """Transforms a function using Haiku modules into a pair of pure functions.
  
      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::
  
          params = init(rng, *a, **k)
          out = apply(params, rng, *a, **k)
  
      Note that the ``rng`` argument is typically not required for ``apply`` and
      passing ``None`` is accepted.
  
      The first thing to do is to define a :class:`Module`. A module encapsulates
      some parameters and a computation on those parameters:
  
      >>> class MyModule(hk.Module):
      ...   def __call__(self, x):
      ...     w = hk.get_parameter("w", [], init=jnp.zeros)
      ...     return x + w
  
      Next, define some function that creates and applies modules. We use
      :func:`transform` to transform that function into a pair of functions that
      allow us to lift all the parameters out of the function (``f.init``) and
      apply the function with a given set of parameters (``f.apply``):
  
      >>> def f(x):
      ...   a = MyModule()
      ...   b = MyModule()
      ...   return a(x) + b(x)
      >>> f = hk.transform(f)
  
      To get the initial state of the module call ``init`` with an example input:
  
      >>> params = f.init(None, 1)
      >>> params
      {'my_module': {'w': ...Array(0., dtype=float32)},
       'my_module_1': {'w': ...Array(0., dtype=float32)}}
  
      You can then apply the function with the given parameters by calling
      ``apply`` (note that since we don't use Haiku's random number APIs to apply
      our network we pass ``None`` as an RNG key):
  
      >>> print(f.apply(params, None, 1))
      2.0
  
      It is expected that your program will at some point produce updated parameters
      and you will want to re-apply ``apply``. You can do this by calling ``apply``
      with different parameters:
  
      >>> new_params = {"my_module": {"w": jnp.array(2.)},
      ...               "my_module_1": {"w": jnp.array(3.)}}
      >>> print(f.apply(new_params, None, 2))
      9.0
  
      If your transformed function needs to maintain internal state (e.g. moving
      averages in batch norm) then see :func:`transform_with_state`.
  
      Args:
        f: A function closing over :class:`Module` instances.
        apply_rng: In the process of being removed. Can only value `True`.
  
      Returns:
        A :class:`Transformed` tuple with ``init`` and ``apply`` pure functions.
      """
      analytics.log_once("transform")
  
      if not apply_rng:
          raise ValueError(
              "The apply_rng argument has been removed and hk.transform "
              "now *always* applies an rng.\n"
              "Replace hk.transform(..., apply_rng=False) with "
              "hk.without_apply_rng(hk.transform(...)).\n"
              "Replace hk.transform(..., apply_rng=True) with hk.transform(...)."
          )
  
      return without_state(transform_with_state(f))
  
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 7 for calling a global function with example usage: transform(f) and returns: without_state