code: def test_run(self, module_fn: ModuleFn, shape, dtype):
      def g(x):
          return module_fn()(x)
  
      f = hk.transform_with_state(g)
  
      def run():
          rng = jax.random.PRNGKey(42)
          x = jnp.zeros(shape, dtype)
          params, state = f.init(rng, x)
          return f.apply(params, state, rng, x)
  
      jax.eval_shape(run)
  
file_path: dm-haiku/haiku/_src/integration/check_tracer_leaks_test.py
imports: from absl.testing import absltest
  from absl.testing import parameterized
  import haiku
  from haiku._src import test_utils
  from haiku._src.integration import descriptors
  import jax
  import jax.numpy
label_desc: Use image node_id 3 for calling the LeakCheckerTest obj's underlying member method code with example usage: obj.test_run(module_fn, shape, dtype) and returns: module_fn, f