code: def _register_shape_inference_hooks(
      module: nn.Module,
      skip_toplevel: bool = True,
      is_leaf: Callable[[nn.Module], bool] | None = None,
  ) -> tuple[
      list[RemovableHandle],
      dict[str, tuple[MutableShape, MutableShape]],
  ]:
      """This method registers every submodule of ``module`` with a hook,
      so that when running forward, the results carry the shape information.
  
      This method only needs to be called once for a module, though it can be called multiple times.
  
      Parameters
      ----------
      module
          The top-level module.
      skip_toplevel
          Whether to wrap the forward of the top-level module.
      is_leaf
          When ``is_leaf`` returns true, its inner submodules will not be wrapped with a hook.
  
      Returns
      -------
      Handles to remove the hooks,
      and a dict that will show the input/output shapes of each submodule after running forward.
      """
      # Handles are used to remove the hooks when exiting.
      handles: list[RemovableHandle] = []
      # Input and output shapes shapes are used to save the shapes of inputs for each submodule.
      # Mapping is from submodule name to input shapes.
      shapes: dict[str, tuple[MutableShape, MutableShape]] = {}
  
      if is_leaf is None:
          is_leaf = is_leaf_module
  
      def save_io(name: str, input: Any, output: Any) -> None:
          shapes[name] = (input, output)
  
      def _iter(name: str, mod: nn.Module) -> None:
          # For self
  
          # Skip the top-level module so that we can at least process keyword arguments for the toppest.
          if not skip_toplevel or name != "":
              # Save input/output shapes for the specific submodule.
              save_io_callback = functools.partial(save_io, name)
              # ... and create the hook.
              hook = functools.partial(
                  module_shape_inference_hook,
                  is_leaf=is_leaf,
                  save_io_callback=save_io_callback,
              )
              # We don't check whether the hook has already been registered here because it's not a big deal,
              # and it will be hard to check whether the hook should be unregistered when exiting.
              handles.append(mod.register_forward_hook(hook))
  
          handles.append(
              mod.register_forward_pre_hook(
                  functools.partial(_current_module_pre_hook, name)
              )
          )
          handles.append(
              mod.register_forward_hook(
                  functools.partial(_current_module_post_hook, name)
              )
          )
  
          # For children
          if is_leaf is None or not is_leaf(mod):
              for child_name, child in mod.named_children():
                  _iter(
                      name + "." + child_name if name else child_name,
                      child,
                  )
  
      _iter("", module)
  
      # shapes contain nothing at this point. It will be filled when forward is called.
      return handles, shapes
  
imports: from __future__ import annotations
  import functools
  import logging
  import traceback
  from typing import Union, Tuple, Callable, Iterable, Any, NamedTuple, Dict, cast, overload
  import torch
  from torch import nn
  from torch.utils.hooks import RemovableHandle
  from nni.mutable import MutableExpression, Mutable, LabeledMutable, Sample, SampleValidationError
  from nni.mutable.mutable import _mutable_equal
  from .misc import is_leaf_module, argument_in_spec
label_desc: Use image node_id 7 for calling a global function with example usage: _register_shape_inference_hooks(module, skip_toplevel, is_leaf) and returns: handles,  shapes