docstr: Test layer_stack with per-layer inputs with multiple arguments.
code: def test_with_per_layer_inputs_multi_args(self):
      """Test layer_stack with per-layer inputs with multiple arguments."""
      width = 4
      batch_size = 5
      stack_height = 3
  
      def f_with_multi_args(x, a, b):
          return (
              basic.Linear(
                  width, w_init=initializers.Constant(jnp.eye(width))
              )(x)
              * a
              + b,
              None,
          )
  
      @multi_transform.without_apply_rng
      @transform.transform
      def hk_fn(x):
          return layer_stack.layer_stack(
              stack_height, with_per_layer_inputs=True
          )(f_with_multi_args)(
              x, jnp.full([stack_height], 2.0), jnp.ones([stack_height])
          )
  
      x = jnp.zeros([batch_size, width])
      key_seq = base.PRNGSequence(19)
      params = hk_fn.init(next(key_seq), x)
      output, z = hk_fn.apply(params, x)
      self.assertIsNone(z)
      self.assertEqual(output.shape, (batch_size, width))
      np.testing.assert_equal(output, np.full([batch_size, width], 7.0))
  
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 11 for calling the LayerStackTest obj's underlying member method code with example usage: obj.test_with_per_layer_inputs_multi_args() and returns: unknown,  None, layer_stack