code: def multi_transform(
      f: Callable[[], tuple[TemplateFn, TreeOfApplyFns]],
  ) -> MultiTransformed:
      """Transforms a collection of functions using Haiku into pure functions.
  
      In many scenarios we have several modules which are used either as primitives
      for several Haiku modules/functions, or whose pure versions are to be reused
      in downstream code. This utility enables this by applying
      :func:`transform` to an arbitrary tree of Haiku functions which share modules
      and have a common ``init`` function.
  
      ``f`` is expected to return a tuple of two elements. First is a ``template``
      Haiku function which provides an example of how all internal Haiku modules are
      connected. This function is used to create a common ``init`` function (with
      your parameters).
  
      The second object is an arbitrary tree of Haiku functions all of which reuse
      the modules connected in the ``template`` function. These functions are
      transformed to pure ``apply`` functions.
  
      Example:
  
      >>> def f():
      ...   encoder = hk.Linear(1, name="encoder")
      ...   decoder = hk.Linear(1, name="decoder")
      ...
      ...   def init(x):
      ...     z = encoder(x)
      ...     return decoder(z)
      ...
      ...   return init, (encoder, decoder)
  
      >>> f = hk.multi_transform(f)
      >>> rng = jax.random.PRNGKey(42)
      >>> x = jnp.ones([1, 1])
      >>> params = f.init(rng, x)
      >>> jax.tree_util.tree_map(jnp.shape, params)
      {'decoder': {'b': (1,), 'w': (1, 1)},
       'encoder': {'b': (1,), 'w': (1, 1)}}
  
      >>> encode, decode = f.apply
      >>> z = encode(params, None, x)
      >>> y = decode(params, None, z)
  
      Args:
        f: A factory function that returns two functions, firstly a common init
          function that creates all modules, and secondly a pytree of apply
          functions which make use of those modules.
  
      Returns:
        A :class:`MultiTransformed` instance which contains a pure init function
          that creates all parameters, and a pytree of pure apply functions that
          given the params apply the given function.
  
      See also:
        :func:`multi_transform_with_state`: Equivalent for modules using state.
      """
      analytics.log_once("multi_transform")
  
      f = multi_transform_with_state(f)
      f = without_state(f)
      return f
  
file_path: dm-haiku/haiku/_src/multi_transform.py
imports: import dataclasses
  import functools
  import inspect
  from typing import Any, Callable, NamedTuple, Optional, TypeVar
  from haiku._src import analytics
  from haiku._src import transform
  from haiku._src import typing
  import jax
label_desc: Use image node_id 2 for calling a global function with example usage: multi_transform(f) and returns: f