code: def test_parameter_in_init(self):
      def f():
          w = base.get_parameter("w", [], init=jnp.zeros)
          s = base.get_state("s", [], init=jnp.zeros)
          init = lambda: None
  
          def add():
              s_add = base.get_state("s", [], init=jnp.zeros)
              w_add = base.get_parameter("w", [], init=jnp.zeros)
              return w, w_add, s, s_add
  
          def sub():
              s_sub = base.get_state("s", [], init=jnp.zeros)
              w_sub = base.get_parameter("w", [], init=jnp.zeros)
              return w, w_sub, s, s_sub
  
          return init, (add, sub)
  
      f = multi_transform.multi_transform_with_state(f)
      params, state = f.init(None)
      self.assertLen(f.apply, 2)
      for apply_fn in f.apply:
          # Check parameter and state reuse inside the transformed function.
          (w, w_apply, s, s_apply), _ = apply_fn(params, state, None)
          self.assertIs(w, w_apply)
          self.assertIs(s, s_apply)
  
imports: import inspect
  from typing import Optional, Union
  from absl.testing import absltest
  from absl.testing import parameterized
  from haiku._src import base
  from haiku._src import multi_transform
  from haiku._src import transform
  from haiku._src import typing
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 3 for calling the MultiTransformTest obj's underlying member method code with example usage: obj.test_parameter_in_init() and returns: init,  add,  sub, w,  w_add,  s,  s_add, w,  w_sub,  s,  s_sub