code: def test_uses_updated_params(self):
      def f():
          mod = transform_flax.lift(Child(), name="foo")
          return mod()
  
      f = transform.transform_with_state(f)
      rng = jax.random.PRNGKey(42)
      params, state = f.init(rng)
      # Modify the parameters from their initial value.
      params, state = jax.tree_map(lambda x: x + 1, (params, state))
      (w_out, flim_out, flam_out), _ = f.apply(params, state, None)
      w = params["foo/~"]["w"]
      flim = state["foo/flim_state/~"]["s"]
      flam = state["foo/flam_state/~"]["s"]
      # We want to assert that the params/state passed in are literally what is
      # returned, this ensures that initialisers did not re-run.
      self.assertIs(w, w_out)
      self.assertIs(flim, flim_out)
      self.assertIs(flam, flam_out)
  
imports: from absl.testing import absltest
  from absl.testing import parameterized
  import flax.errors
  import flax.linen
  from haiku._src import base
  from haiku._src import module
  from haiku._src import transform
  from haiku._src.flax import transform_flax
  import jax
  import jax.numpy
label_desc: Use image node_id 7 for calling the TransformFlaxTest obj's underlying member method code with example usage: obj.test_uses_updated_params() and returns: mod