code: def repeat(layer: Layer, n: int) -> InternalLayer:
      """Compose `layer` in a compiled loop `n` times.
  
      Equivalent to `serial(*([layer] * n))`, but allows faster compilation time
      for large `n` (but same runtime).
  
      .. warning::
        `apply_fn` of the `layer` is assumed to keep the activation (`x`) shape
        unchanged.
  
      .. warning::
        `kernel_fn` of the `layer` is assumed to keep the
        :class:`~neural_tangents.Kernel` metadata unchanged. This is most notably
        not satisfied in :obj:`~neural_tangents.stax.Conv` and other convolutional
        layers which flip the `is_reversed` attribute with each application. A
        workaround is to either use `serial(*([layer] * n))`, or to use
        `repeat(serial(layer, layer), n // 2)` instead of `repeat(layer, n)` for an
        even `n`, i.e. to use two (or, generally, any even number of) convolutions
        per `layer` instead of one (or, generally, any odd number), such that
        `layer` does not alter the `is_reversed` attribute. Similar caution should
        be applied to other :class:`~neural_tangents.Kernel` attributes.
  
      See Also:
        `RepeatTest` in `tests/stax/combinators_test.py` for examples and
        :obj:`~neural_tangents.stax.serial` for unrolled composition.
  
      Example:
        >>> from neural_tangents import stax
        >>> #
        >>> layer = stax.serial(stax.Dense(128), stax.Relu())
        >>> depth = 100
        >>> #
        >>> # Unrolled loop:
        >>> nn_unrolled = stax.serial(*([layer] * depth))
        >>> #
        >>> # Compiled loop:
        >>> nn_compiled = stax.repeat(layer, depth)
        >>> # `nn_unrolled` and `nn_compiled` perform the same computation, but
        >>> # `nn_compiled` compiles faster and with smaller memory footprint.
  
      Args:
        layer:
          layer to be repeated. Outputs must have the same shape and other metadata
          as inputs.
  
        n:
          number of times to repeat a layer (depth).
  
      Returns:
        A new layer, meaning an `(init_fn, apply_fn, kernel_fn)` triple,
        representing the repeated composition of `layer` `n` times.
      """
      init_fn, apply_fn, kernel_fn = layer
  
      def init_fn_repeat(rng, input_shape):
          out_shape, _ = init_fn(rng, input_shape)
          if out_shape != input_shape:
              raise ValueError(
                  f"`init_fn` produces a different output shape {out_shape} than the "
                  f"input shape {input_shape}. Please use the `serial(*([layer] * n)`) "
                  f"construction in this setting."
              )
  
          def init_fn_scan(rng, params):
              rng, layer_rng = random.split(rng)
              out_shape, params = init_fn(layer_rng, input_shape)
              return rng, params
  
          _, params = lax.scan(init_fn_scan, rng, None, n)
          return out_shape, params
  
      def apply_fn_repeat(params, inputs, **kwargs):
          def apply_fn_scan(x, params):
              return apply_fn(params, x, **kwargs), None
  
          outputs, _ = lax.scan(apply_fn_scan, inputs, params, n)
          return outputs
  
      @requires(**get_req(kernel_fn))
      def kernel_fn_repeat(
          k: NTTree[Kernel], **kwargs
      ) -> NTTree[Kernel]:
          if n > 0:
              k = kernel_fn(k, **kwargs)
  
              def kernel_fn_scan(k, _):
                  k = kernel_fn(k, **kwargs)
                  return k, None
  
              k, _ = lax.scan(kernel_fn_scan, k, None, n - 1)
  
          return k
  
      return init_fn_repeat, apply_fn_repeat, kernel_fn_repeat
  
imports: import operator
  from typing import Any, Callable
  import warnings
  import frozendict
  import jax
  from jax import lax
  from jax import random
  import jax.example_libraries.stax
  from ..utils.kernel import Kernel
  from ..utils.typing import InternalLayer
  from ..utils.typing import Layer
  from ..utils.typing import LayerKernelFn
  from ..utils.typing import NTTree
  from ..utils.typing import NTTrees
  from ..utils.typing import Shapes
  from .requirements import Diagonal
  from .requirements import get_req
  from .requirements import layer
  from .requirements import requires
label_desc: Use image node_id 2 for calling a global function with example usage: repeat(layer, n) and returns: init_fn_repeat,  apply_fn_repeat,  kernel_fn_repeat, out_shape,  params, outputs, k, rng,  params, apply_fn,  None, k,  None