code: def multi_transform_with_state(
      f: Callable[[], tuple[TemplateFn, TreeOfApplyFns]],
  ) -> MultiTransformedWithState:
      """Transforms a collection of functions using Haiku into pure functions.
  
      See :func:`multi_transform` for more details.
  
      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_with_state(f)
      >>> rng = jax.random.PRNGKey(42)
      >>> x = jnp.ones([1, 1])
      >>> params, state = 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, state = encode(params, state, None, x)
      >>> y, state = decode(params, state, None, z)
  
      Args:
        f: Function returning a "template" function and an arbitrary
          tree of functions using modules connected in the template function.
  
      Returns:
        An ``init`` function and a tree of pure ``apply`` functions.
  
      See also:
        - :func:`transform_with_state`: Transform a single apply function.
        - :func:`multi_transform`: Transform multiple apply functions without state.
      """
      analytics.log_once("multi_transform_with_state")
  
      def init_fn(
          *args, **kwargs
      ) -> tuple[hk.MutableParams, hk.MutableState]:
          """Returns initial state for the transformed functions."""
          return f()[0](*args, **kwargs)
  
      init_fn = hk.transform_with_state(init_fn).init
  
      def apply_fn_i(i):
          def apply_fn(*args, **kwargs):
              """Applies the transformed function at the given inputs."""
              return jax.tree_util.tree_leaves(f()[1])[i](
                  *args, **kwargs
              )
  
          return apply_fn
  
      # We need to find out the structure of f()[1], including how many
      # functions there are, so that we can transform them individually and repack
      # into the same tree structure. It's valid for modules to declare parameters
      # in their constructor, so we need to create something that looks like
      # hk.Params in order to do this. `jax.eval_shape` interprets the function
      # abstractly, ie no real params are created, and we don't need to touch the
      # accelerator. This means hardcoding the RNG below is fine.
      def get_output_treedef() -> Box:
          rng = jax.random.PRNGKey(42)  # This is fine, see above
          fns = hk.transform_with_state(lambda: f()[1])
          apply_fns, _ = fns.apply(*fns.init(rng), rng)
          return Box(jax.tree_util.tree_structure(apply_fns))
  
      output_treedef = jax.eval_shape(get_output_treedef).python_value
      apply_fns = make_tree(
          lambda i: hk.transform_with_state(apply_fn_i(i)).apply,
          output_treedef,
      )
  
      return MultiTransformedWithState(init_fn, apply_fns)
  
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 1 for calling a global function with example usage: multi_transform_with_state(f) and returns: MultiTransformedWithState, f, apply_fn, Box, jax