code: def test_with_container_state(self):
      width = 2
      batch_size = 2
      stack_height = 3
  
      def f_with_container_state(x):
          hk_layer = basic.Linear(
              width, w_init=initializers.Constant(jnp.eye(width))
          )
          layer_output = hk_layer(x)
          layer_state = {
              "raw_output": layer_output,
              "output_projection": jnp.sum(layer_output),
          }
          return layer_output + jnp.ones_like(layer_output), layer_state
  
      @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_container_state)(x)
  
      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.assertEqual(
          z["raw_output"].shape, (stack_height, batch_size, width)
      )
      self.assertEqual(output.shape, (batch_size, width))
      self.assertEqual(z["output_projection"].shape, (stack_height,))
      np.testing.assert_equal(
          np.sum(z["output_projection"]), np.array(12.0)
      )
      np.testing.assert_equal(
          np.all(
              z["raw_output"]
              == np.array([0.0, 1.0, 2.0])[..., None, None]
          ),
          np.array(True),
      )
  
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 12 for calling the LayerStackTest obj's underlying member method code with example usage: obj.test_with_container_state() and returns: unknown,  layer_state, layer_stack