code: def module_type(module_fn: ModuleFn) -> type[hk.Module]:
      f = hk.transform_with_state(lambda: type(unwrap(module_fn())))
      return f.apply(*f.init(jax.random.PRNGKey(42)), None)[0]
  
imports: from collections.abc import Sequence
  from typing import Any, Callable, NamedTuple
  import haiku
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 4 for calling a global function with example usage: module_type(module_fn) and returns: f