code: def get_module_cls(module_fn: ModuleFn) -> type[hk.Module]:
      get_cls = lambda: type(descriptors.unwrap(module_fn()))
      return hk.testing.transform_and_run(get_cls)()
  
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 1 for calling a global function with example usage: get_module_cls(module_fn) and returns: hk