code: def _name_scopes_to_modules(module: Module):
      """Converts name scopes to nested Modules, as if a call jaxpr were present."""
      expressions = list(module.expressions)
      del module.expressions[:]
  
      if module.flops:
          # Flops are recomputed below.
          module.flops = 0
  
      # The nested set of Modules corresponding to the current name scope stack.
      module_stack = [module]
  
      for e in expressions:
          if e.submodule is not None:
              _name_scopes_to_modules(e.submodule)
  
          # Close open scopes.
          while len(module_stack) > len(e.name_stack) + 1:
              module_stack.pop()
  
          for i, n in enumerate(e.name_stack):
              i = i + 1  # Offset to account for the outermost module.
              if i < len(module_stack) and n != module_stack[i].name:
                  # Entering a different name scope.
                  while len(module_stack) > i:
                      module_stack.pop()
  
              if i == len(module_stack):
                  # Represent the scope as a dummy module.
                  scope = Expression(
                      primitive="name_scope", invars="", outvars=""
                  )
                  scope.submodule = Module(name=n)
                  module_stack[-1].expressions.append(scope)
                  module_stack.append(scope.submodule)
  
          # Expressions belong to the corresponding innermost module.
          module_stack[-1].expressions.append(e)
          if e.flops:
              for m in module_stack:
                  m.flops = (m.flops or 0) + e.flops
  
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 2 for calling a global function with example usage: _name_scopes_to_modules(module) without return types