code: def _add_param_counts(
      model_info: Module,
      prefix: str,
      invocations: Mapping[str, summarise.MethodInvocation],
  ):
      """Adds parameter count to the module info, if present."""
      if prefix in invocations:
          invocation = invocations[prefix]
          model_info.total_param_size = sum(
              spec.size
              for spec in invocation.module_details.params.values()
          )
          model_info.total_state_size = sum(
              spec.size
              for spec in invocation.module_details.state.values()
          )
          for k, s in invocation.module_details.params.items():
              model_info.param_info[
                  k.replace("/~/", "/").replace(prefix, ".")
              ] = str(s)
          for k, s in invocation.module_details.state.items():
              model_info.state_info[
                  k.replace("/~/", "/").replace(prefix, ".")
              ] = str(s)
  
      for expr in model_info.expressions:
          submodule = expr.submodule
          if submodule:
              _add_param_counts(
                  submodule,
                  os.path.join(prefix, submodule.name),
                  invocations,
              )
  
imports: from collections.abc import Mapping, Sequence
  import dataclasses
  import itertools
  import logging
  import os
  import sys
  from typing import Any, Callable, NamedTuple, Optional
  from haiku._src import summarise
  import jax
  import jax.core
label_desc: Use image node_id 9 for calling a global function with example usage: _add_param_counts(model_info, prefix, invocations) without return types