code: def test_reverse(self):
      # The layer stack below runs iteratively the update equation:
      # x_n = n * alpha * (x_{n-1} + 1)
      # with x_0 = 1, for n={1, ..., N}, where N = stack_height
      # The reverse layer stack as a result runs the update equation:
      # y_{n-1} = (N - n + 1) * alpha * (y_n + 1)
      # with y_N = 1, for n={N-1, ..., 0}, where N = stack_height
      width = 2
      batch_size = 3
      stack_height = 4
      alpha = jnp.power(24, -1.0 / stack_height)
      forward, backward = self._compute_weights(stack_height, alpha)
  
      def inner_fn(x):
          # Here we initialize the layer to an identity + 1, while later we multiply
          # each parameter by the index `n`.
          return basic.Linear(
              x.shape[1],
              w_init=initializers.Constant(jnp.eye(x.shape[1])),
              b_init=initializers.Constant(1.0),
          )(x)
  
      @multi_transform.without_apply_rng
      @transform.transform
      def hk_fn(x, reverse=False):
          return layer_stack.layer_stack(stack_height)(inner_fn)(
              x, reverse=reverse
          )
  
      key_seq = base.PRNGSequence(19)
      init_value = 1 + jax.random.uniform(
          next(key_seq), [batch_size, width]
      )
  
      def mul_by_m(x):
          m_x = jnp.arange(stack_height) + 1
          while m_x.ndim < x.ndim:
              m_x = m_x[..., None]
          return x * m_x * alpha
  
      params = jax.tree_util.tree_map(
          mul_by_m, hk_fn.init(next(key_seq), init_value)
      )
  
      a, b = forward[-1]
      x_n = hk_fn.apply(params, init_value)
      np.testing.assert_allclose(x_n, a * init_value + b, rtol=1e-6)
  
      a, b = backward[-1]
      y_0 = hk_fn.apply(params, init_value, reverse=True)
      np.testing.assert_allclose(y_0, a * init_value + b, rtol=1e-6)
  
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 14 for calling the LayerStackTest obj's underlying member method code with example usage: obj.test_reverse() and returns: basic, layer_stack, unknown