code: def test_nested_stacks(self):
      def stack_fn(x):
          def layer_fn(x):
              return basic.Linear(100)(x)
  
          outer_fn = layer_stack.layer_stack(10)(layer_fn)
  
          layer_outer = layer_stack.layer_stack(20)(outer_fn)
          return layer_outer(x)
  
      hk_mod = transform.transform(stack_fn)
      apply_rng, init_rng = jax.random.split(jax.random.PRNGKey(0))
  
      params = hk_mod.init(init_rng, jnp.zeros([10, 100]))
  
      hk_mod.apply(params, apply_rng, jnp.zeros([10, 100]))
  
      (p,) = params.values()
  
      assert p["w"].shape == (20, 10, 100, 100), p["w"].shape
      assert p["b"].shape == (20, 10, 100), p["b"].shape
  
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 10 for calling the LayerStackTest obj's underlying member method code with example usage: obj.test_nested_stacks() and returns: layer_outer, basic