code: def summarize(d: descriptors.ModuleDescriptor) -> Mapping[str, Any]:
      """Generates a summary of the given descriptor."""
      f = hk.transform_with_state(
          lambda x: d.create()(x)
      )  # pylint: disable=unnecessary-lambda
      x = jnp.ones(d.shape, d.dtype)
      rng = jax.random.PRNGKey(42)
      params, state = map(
          hk.data_structures.to_mutable_dict,
          jax.eval_shape(f.init, rng, x),
      )
      out = {"module": module_name(d), "input": format_tensor(x)}
      if params:
          out["param_size"] = int(hk.data_structures.tree_size(params))
          out["param_bytes"] = int(
              hk.data_structures.tree_bytes(params)
          )
          out["params"] = jax.tree_util.tree_map(format_tensor, params)
      if state:
          out["state_size"] = int(hk.data_structures.tree_size(state))
          out["state_bytes"] = int(hk.data_structures.tree_bytes(state))
          out["state"] = jax.tree_util.tree_map(format_tensor, state)
      return out
  
file_path: dm-haiku/haiku/_src/integration/checkpoint_utils.py
imports: from collections.abc import Mapping
  from typing import Any
  import haiku
  from haiku._src import utils
  from haiku._src.integration import descriptors
  import jax
  import jax.numpy
label_desc: Use image node_id 3 for calling a global function with example usage: summarize(d) and returns: out