# Path: alphafold/model/layer_stack.py
# def layer_stack(num_layers: int,
#                 with_state=False,
#                 unroll: int = 1,
#                 name: Optional[str] = None):
#   """Utility to wrap a Haiku function and recursively apply it to an input.
# 
#   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 `jnp.ndarray` at the leaf nodes. Note
#   that kwargs are not supported, neither are functions with variable number
#   of parameters (specified by `*args`).
# 
#   If `with_state=False` then the new, wrapped function can be understood as
#   performing the following:
#   ```
#   for i in range(num_layers):
#     x = f(x)
#   return x
#   ```
# 
#   And if `with_state=True`, assuming `f` takes two arguments on top of `x`:
#   ```
#   for i in range(num_layers):
#     x, zs[i] = f(x, ys_0[i], ys_1[i])
#   return x, zs
#   ```
#   The code using `layer_stack` for the above function would be:
#   ```
#   def f(x, y_0, y_1):
#     ...
#     return new_x, z
#   x, zs = layer_stack.layer_stack(num_layers,
#                                   with_state=True)(f)(x, ys_0, ys_1)
#   ```
# 
#   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_state: Whether or not to pass per-layer state to the wrapped function.
#     unroll: the unroll used by `scan`.
#     name: Name of the Haiku context.
# 
#   Returns:
#     Callable that will produce a layer stack when called with a valid function.
#   """
#   def iterate(f):
#     if with_state:
#       @functools.wraps(f)
#       def wrapped(x, *args):
#         for ys in args:
#           assert ys.shape[0] == num_layers
#         return _LayerStackWithState(
#             f, num_layers, unroll=unroll, name=name)(x, *args)
#     else:
#       _check_no_varargs(f)
#       @functools.wraps(f)
#       def wrapped(*args):
#         ret = _LayerStackNoState(
#             f, num_layers, unroll=unroll, name=name)(args, None)[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


import functools
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
import scipy
from absl.testing import absltest
from absl.testing import parameterized
from alphafold.model import layer_stack
          sliced_layers_params[k_new] = {}
        sliced_layers_params[k_new][inner_k] = var_slice
  return sliced_layers_params


class LayerStackTest(parameterized.TestCase):

  @parameterized.parameters([1, 2, 4])
  def test_layer_stack(self, unroll):
    """Compare layer_stack to the equivalent unrolled stack.

    Tests that the layer_stack application of a Haiku layer function is
    equivalent to repeatedly applying the layer function in an unrolled loop.

    Args:
      unroll: Number of unrolled layers.
    """
    num_layers = 20

    def inner_fn(x):
      x += hk.Linear(100, name='linear1')(x)
      x += hk.Linear(100, name='linear2')(x)
      return x

    def outer_fn_unrolled(x):
      for _ in range(num_layers):
        x = inner_fn(x)
      return x

    def outer_fn_layer_stack(x):

# Path: alphafold/model/layer_stack.py
# def layer_stack(num_layers: int,
#                 with_state=False,
#                 unroll: int = 1,
#                 name: Optional[str] = None):
#   """Utility to wrap a Haiku function and recursively apply it to an input.
# 
#   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 `jnp.ndarray` at the leaf nodes. Note
#   that kwargs are not supported, neither are functions with variable number
#   of parameters (specified by `*args`).
# 
#   If `with_state=False` then the new, wrapped function can be understood as
#   performing the following:
#   ```
#   for i in range(num_layers):
#     x = f(x)
#   return x
#   ```
# 
#   And if `with_state=True`, assuming `f` takes two arguments on top of `x`:
#   ```
#   for i in range(num_layers):
#     x, zs[i] = f(x, ys_0[i], ys_1[i])
#   return x, zs
#   ```
#   The code using `layer_stack` for the above function would be:
#   ```
#   def f(x, y_0, y_1):
#     ...
#     return new_x, z
#   x, zs = layer_stack.layer_stack(num_layers,
#                                   with_state=True)(f)(x, ys_0, ys_1)
#   ```
# 
#   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_state: Whether or not to pass per-layer state to the wrapped function.
#     unroll: the unroll used by `scan`.
#     name: Name of the Haiku context.
# 
#   Returns:
#     Callable that will produce a layer stack when called with a valid function.
#   """
#   def iterate(f):
#     if with_state:
#       @functools.wraps(f)
#       def wrapped(x, *args):
#         for ys in args:
#           assert ys.shape[0] == num_layers
#         return _LayerStackWithState(
#             f, num_layers, unroll=unroll, name=name)(x, *args)
#     else:
#       _check_no_varargs(f)
#       @functools.wraps(f)
#       def wrapped(*args):
#         ret = _LayerStackNoState(
#             f, num_layers, unroll=unroll, name=name)(args, None)[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


# Path: alphafold/model/layer_stack_test.py
import functools
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
import scipy
from absl.testing import absltest
from absl.testing import parameterized
from alphafold.model import layer_stack

          sliced_layers_params[k_new] = {}
        sliced_layers_params[k_new][inner_k] = var_slice
  return sliced_layers_params


class LayerStackTest(parameterized.TestCase):

  @parameterized.parameters([1, 2, 4])
  def test_layer_stack(self, unroll):
    """Compare layer_stack to the equivalent unrolled stack.

    Tests that the layer_stack application of a Haiku layer function is
    equivalent to repeatedly applying the layer function in an unrolled loop.

    Args:
      unroll: Number of unrolled layers.
    """
    num_layers = 20

    def inner_fn(x):
      x += hk.Linear(100, name='linear1')(x)
      x += hk.Linear(100, name='linear2')(x)
      return x

    def outer_fn_unrolled(x):
      for _ in range(num_layers):
        x = inner_fn(x)
      return x

    def outer_fn_layer_stack(x):

      stack = layer_stack.layer_stack(num_layers, unroll=unroll)(inner_fn)