code: def test_transform_with_state_not_mutating_state_after_ini(self):
      def f():
          s1 = base.get_state("s1", [], init=jnp.zeros)
          s2 = base.get_state("s2", [], init=jnp.ones)
          return s1, s2
  
      f = transform.transform_with_state(f)
      m = flax_module.Module(f)
      rng = jax.random.PRNGKey(42)
      variables = m.init(rng)
      self.assertEqual(variables, {"state": {"~": {"s1": 0, "s2": 1}}})
  
      # NOTE: Intentionally not making state collection mutable.
      s1, s2 = m.apply(variables)
      self.assertIs(s1, variables["state"]["~"]["s1"])
      self.assertIs(s2, variables["state"]["~"]["s2"])
  
imports: from absl.testing import absltest
  from absl.testing import parameterized
  import flax.linen
  from haiku._src import base
  from haiku._src import filtering
  from haiku._src import module
  from haiku._src import transform
  from haiku._src.flax import flax_module
  from haiku._src.flax import utils
  from haiku._src.nets import resnet
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 3 for calling the FlaxModuleTest obj's underlying member method code with example usage: obj.test_transform_with_state_not_mutating_state_after_ini() and returns: s1,  s2