code: def layer_stack(
      num_layers: int,
      with_per_layer_inputs=False,
      unroll: int = 1,
      pass_reverse_to_layer_fn: bool = False,
      transparent: bool = False,
      transparency_map: Optional[LayerStackTransparencyMapping] = None,
      name: Optional[str] = None,
  ):
      """Utility to wrap a Haiku function and recursively apply it to an input.
  
      This can be used to improve model compile times.
  
      A function is valid if it uses only explicit position parameters, and
      its return type matches its input type. The position parameters can be
      arbitrarily nested structures with ``jax.Array`` at the leaf nodes. Note
      that kwargs are not supported, neither are functions with variable number
      of parameters (specified by ``*args``).
  
      Note that `layer_stack` cannot at the moment be used with functions that build
      Haiku modules with state.
  
      If ``with_per_layer_inputs=False`` then the new, wrapped function can be
      understood as performing the following:
  
      >>> f = lambda x: x+1
      >>> num_layers = 4
      >>> x = 0
      >>> for i in range(num_layers):
      ...   x = f(x)
      >>> x
      4
  
      And if ``with_per_layer_inputs=True``, assuming ``f`` takes two arguments on
      top of ``x``:
  
      >>> f = lambda x, y0, y1: (x+1, y0+y1)
      >>> num_layers = 4
      >>> x = 0
      >>> ys_0 = [1, 2, 3, 4]
      >>> ys_1 = [5, 6, 7, 8]
      >>> zs = []
      >>> for i in range(num_layers):
      ...   x, z = f(x, ys_0[i], ys_1[i])
      ...   zs.append(z)
      >>> x, zs
      (4, [6, 8, 10, 12])
  
      The code using ``layer_stack`` for the above function would be:
  
      >>> f = lambda x, y0, y1: (x+1, y0+y1)
      >>> num_layers = 4
      >>> x = 0
      >>> ys_0 = jnp.array([1, 2, 3, 4])
      >>> ys_1 = jnp.array([5, 6, 7, 8])
      >>> stack = hk.layer_stack(num_layers, with_per_layer_inputs=True)
      >>> x, zs = stack(f)(x, ys_0, ys_1)
      >>> print(x, zs)
      4 [ 6  8 10 12]
  
      Check the tests in ``layer_stack_test.py`` for further examples.
  
      Crucially, any parameters created inside ``f`` will not be shared across
      iterations.
  
      Args:
        num_layers: The number of times to iterate the wrapped function.
        with_per_layer_inputs: Whether or not to pass per-layer inputs to the
          wrapped function.
        unroll: the unroll used by ``scan``.
        pass_reverse_to_layer_fn: Whether or not to pass the ``reverse`` keyword to
          the function ``f``, so that it is aware if the layer stack is being run
          forward or in reverse (and the underlying ``scan``). To run the layer
          stack in reverse you need to pass in ``reverse=True`` to the call to the
          layer stack.
        transparent: Whether to apply layer_stack transparently.  When this is True,
          and a correct transparency_map is provided, the parameters are generated
          in such a way that layer_stack can be replaced by a regular for loop
          without changing the parameter tree.
        transparency_map: How to map stacked module names to flat names and reverse.
          See ``LayerStackTransparencyMapping`` and ``layer_stack_test.py`` for an
          example.
        name: name of the Haiku context.
  
      Returns:
        Callable that will produce a layer stack when called with a valid function.
      """
      if transparent and transparency_map is None:
          raise ValueError(
              "transparency_map must be provided with transparent=True."
          )
  
      if not name:
          if with_per_layer_inputs:
              name = "__layer_stack_with_per_layer"
          else:
              name = "__layer_stack_no_per_layer"
  
      def iterate(f):
          if with_per_layer_inputs:
  
              @functools.wraps(f)
              def wrapped(x, *args, **kwargs):
                  for ys in jax.tree_util.tree_leaves(args):
                      assert (
                          ys.shape[0] == num_layers
                      ), f"{ys.shape[0]} != {num_layers}"
                  mod = _LayerStackWithPerLayer(
                      f,
                      num_layers,
                      unroll=unroll,
                      pass_reverse_to_layer_fn=pass_reverse_to_layer_fn,
                      transparency_map=transparency_map,
                      name=name,
                  )
                  return mod(x, *args, **kwargs)
  
          else:
              _check_no_varargs(f)
  
              @functools.wraps(f)
              def wrapped(*args, **kwargs):
                  mod = _LayerStackNoPerLayer(
                      f,
                      num_layers,
                      unroll=unroll,
                      pass_reverse_to_layer_fn=pass_reverse_to_layer_fn,
                      transparency_map=transparency_map,
                      name=name,
                  )
                  ret = mod(x=args, **kwargs)[0]
                  if len(args) == 1:
                      # If the function takes a single argument, we must also return a
                      # single value, and not a tuple of length 1.
                      ret = ret[0]
                  return ret
  
          return wrapped
  
      return iterate
  
imports: import collections
  import functools
  import inspect
  from typing import Any, Callable, Optional, Protocol, Union
  from haiku._src import base
  from haiku._src import lift
  from haiku._src import module
  from haiku._src import transform
  import jax
  import jax.numpy
label_desc: Use image node_id 5 for calling a global function with example usage: layer_stack(num_layers, with_per_layer_inputs, unroll, pass_reverse_to_layer_fn, transparent, transparency_map, name) and returns: iterate, wrapped, mod, ret