code: def submodule_input_output_shapes(
      model: nn.Module,
      *args: ShapeTensor,
      is_leaf: Callable[[nn.Module], bool] | None = None,
      **kwargs: ShapeTensor,
  ) -> dict[str, tuple[MutableShape, MutableShape]]:
      """Get the dict of all the symbolic shapes of the inputs and outputs of all the submodules.
  
      Parameters
      ----------
      model
          The top-level module.
      *args
          Positional arguments to the model's forward.
      **kwargs
          Keyword arguments.
      is_leaf
          A function that takes a module and returns whether the module is a leaf module.
          The submodule of a leaf module will not appear in the result dict,
          which means we won't complain about their shapes even if missing (though there might be still warnings).
          By default, a module is a leaf module if it has no children or it's defined in torch.nn or its parameterized version.
          See :func:`~nni.nas.pytorch.utils.is_leaf_module` for details.
      """
      handles, shapes = _register_shape_inference_hooks(
          model, skip_toplevel=False, is_leaf=is_leaf
      )
  
      with torch.no_grad():
          model(*args, **kwargs)
  
      for handle in handles:
          handle.remove()
  
      return 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 4 for calling a global function with example usage: submodule_input_output_shapes(model) and returns: shapes