docstr: Test @layer_stack when the function gets per-layer inputs.
code: def test_threading(self):
      """Test @layer_stack when the function gets per-layer inputs."""
      n = 5
  
      @layer_stack.layer_stack(n, with_per_layer_inputs=True)
      def f(x, y):
          x = x + y * jax.nn.one_hot(y, len(x)) / 10
          return x, 2 * y
  
      @multi_transform.without_apply_rng
      @transform.transform
      def g(x, ys):
          x, zs = f(x, ys)
          # Check here to catch issues at init time
          self.assertEqual(zs.shape, (n,))
          return x, zs
  
      rng = jax.random.PRNGKey(7)
      x = np.zeros(n)
      ys = np.arange(n).astype(np.float32)
      params = g.init(rng, x, ys)
      x, zs = g.apply(params, x, ys)
      self.assertTrue(np.allclose(x, [0, 0.1, 0.2, 0.3, 0.4]))
      self.assertTrue(np.all(zs == 2 * ys))
  
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 9 for calling the LayerStackTest obj's underlying member method code with example usage: obj.test_threading() and returns: x,  unknown, x,  zs