docstr: Compare layers_stack to the equivalent unrolled stack.
  
  Similar to `test_layer_stack`, but use a function that takes more than one
  argument.
code: def test_layer_stack_multi_args(self):
      """Compare layers_stack to the equivalent unrolled stack.
  
      Similar to `test_layer_stack`, but use a function that takes more than one
      argument.
      """
      num_layers = 20
  
      def inner_fn(x, y):
          x_out = x + basic.Linear(100, name="linear1")(y)
          y_out = y + basic.Linear(100, name="linear2")(x)
          return x_out, y_out
  
      def outer_fn_unrolled(x, y):
          for _ in range(num_layers):
              x, y = inner_fn(x, y)
          return x, y
  
      def outer_fn_layer_stack(x, y):
          stack = layer_stack.layer_stack(num_layers)(inner_fn)
          return stack(x, y)
  
      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])
      y = jax.random.uniform(jax.random.PRNGKey(1), [10, 256, 100])
  
      rng_init = jax.random.PRNGKey(42)
  
      params = layer_stack_fn.init(rng_init, x, y)
  
      sliced_params = _slice_layers_params(params)
  
      unrolled_x, unrolled_y = unrolled_fn.apply(
          sliced_params, None, x, y
      )
      layer_stack_x, layer_stack_y = layer_stack_fn.apply(
          params, None, x, y
      )
  
      np.testing.assert_allclose(unrolled_x, layer_stack_x, atol=1e-3)
      np.testing.assert_allclose(unrolled_y, layer_stack_y, atol=1e-3)
  
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 4 for calling the LayerStackTest obj's underlying member method code with example usage: obj.test_layer_stack_multi_args() and returns: x_out,  y_out, x,  y, stack