docstr: Compare layers_stack gradients to the equivalent unrolled stack.
  
  Tests that the layers_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.
code: def test_layer_stack_grads(self, unroll):
      """Compare layers_stack gradients to the equivalent unrolled stack.
  
      Tests that the layers_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 += basic.Linear(100, name="linear1")(x)
          x += basic.Linear(100, name="linear2")(x)
          x /= jnp.mean(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
          )
          return stack(x)
  
      unrolled_fn = transform.transform(outer_fn_unrolled)
      layer_stack_fn = transform.transform(outer_fn_layer_stack)
  
      x = jax.random.uniform(jax.random.PRNGKey(0), [10, 256, 100])
  
      rng_init = jax.random.PRNGKey(42)
  
      params = layer_stack_fn.init(rng_init, x)
  
      sliced_params = _slice_layers_params(params)
  
      unrolled_grad = jax.grad(
          lambda p, x: jnp.mean(unrolled_fn.apply(p, None, x))
      )(sliced_params, x)
      layer_stack_grad = jax.grad(
          lambda p, x: jnp.mean(layer_stack_fn.apply(p, None, x))
      )(params, x)
  
      assert_fn = functools.partial(
          np.testing.assert_allclose, atol=1e-4, rtol=1e-4
      )
  
      jax.tree_util.tree_map(
          assert_fn,
          unrolled_grad,
          _slice_layers_params(layer_stack_grad),
      )
  
imports: import functools
  import re
  from typing import Optional
  from absl.testing import absltest
  from absl.testing import parameterized
  from haiku._src import base
  from haiku._src import basic
  from haiku._src import config
  from haiku._src import initializers
  from haiku._src import layer_stack
  from haiku._src import module
  from haiku._src import multi_transform
  from haiku._src import transform
  from haiku._src import utils
  import jax
  import jax.numpy
  import numpy
  from scipy import stats
label_desc: Use image node_id 7 for calling the LayerStackTest obj's underlying member method code with example usage: obj.test_layer_stack_grads(unroll) and returns: x, x, stack