code: def make_model_info(
      f: Callable[..., Any],
      name: Optional[str] = None,
      include_module_info: bool = True,
      compute_flops: Optional[ComputeFlopsFn] = None,
      axis_env: Optional[Sequence[tuple[Any, int]]] = None,
  ) -> Callable[..., Module]:
      """Creates a function that computes flop, param and state information.
  
      Args:
        f: The function for which to compute information. Haiku modules and
          jax.named_call expressions will be represented as nested Modules in the
          result.
        name: Optional, the name of the root expression.
        include_module_info: Whether to include parameter and state count
          information for haiku modules. Can be slow for very large computations.
        compute_flops: Optional, a function that returns an estimate of the number
          of flops required to execute an equation.
        axis_env: Sizes of pmapped axes.  See docs of jax.make_jaxpr for details.
  
      Returns:
        A wrapped version of `f` that when applied to example arguments returns a
        `Module` representation of `f` for those arguments.
  
      `Module` and `Expression` contain high level information about JAX operations
      (jaxprs) and can be visualized in concise and interactive formats; see
      `format_module`, `as_html_page` or `as_html`.
      """
      if not name:
          name = f.__name__
      make_jaxpr = jax.make_jaxpr(f, axis_env=axis_env)
      if include_module_info:
          # Wrap f in a lambda so eval_summary doesn't try to un-transform it.
          # TODO(tomhennigan): remove lambda trick
          make_module_info = summarise.eval_summary(
              lambda *a, **k: f(*a, **k)
          )  # pylint: disable=unnecessary-lambda
  
      def make_module(*args, **kwargs):
          old_limit = sys.getrecursionlimit()
  
          try:
              # Increase recursion limit as graphs may be very deep
              sys.setrecursionlimit(int(10e3))
  
              jaxpr = make_jaxpr(*args, **kwargs).jaxpr
  
              # Compute flops for all expressions.
              module = Module(name=name)
              _process_jaxpr(
                  jaxpr,
                  compute_flops,
                  scope=_ModuleScope(named_call_id="0"),
                  seen=set(),
                  module=module,
              )
  
              _name_scopes_to_modules(module)
  
              if include_module_info:
                  # Add haiku param and state counts for all haiku modules.
                  module_infos = make_module_info(*args, **kwargs)
                  by_name = {
                      i.module_details.module.module_name: i
                      for i in module_infos
                  }
                  by_name = {
                      k.replace("/~/", "/"): v
                      for k, v in by_name.items()
                  }
  
                  for expr in module.expressions:
                      submodule = expr.submodule
                      if submodule:
                          _add_param_counts(
                              submodule, submodule.name, by_name
                          )
          finally:
              sys.setrecursionlimit(old_limit)
          return module
  
      return make_module
  
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 1 for calling a global function with example usage: make_model_info(f, name, include_module_info, compute_flops, axis_env) and returns: make_module, module