code: def shape_inference(
      module: nn.Module,
      *args: ShapeTensor,
      is_leaf: Callable[[nn.Module], bool] | None = None,
      **kwargs: ShapeTensor,
  ) -> ShapeTensor | tuple[ShapeTensor, ...]:
      """
      Running the forward of module with :class:`ShapeTensor` as inputs,
      and return the output :class:`ShapeTensor`.
  
      Parameters
      ----------
      module
          The module to run.
      *args
          The inputs to the module.
      **kwargs
          The keyword arguments to the module.
      is_leaf
          A function that takes a module and returns whether the module is a leaf module.
  
      See Also
      --------
      submodule_input_output_shapes
      """
      handles, _ = _register_shape_inference_hooks(
          module, is_leaf=is_leaf
      )
  
      with torch.no_grad():
          outputs = module(*args, **kwargs)
      result = _module_shape_inference_impl(
          module, outputs, *args, **kwargs
      )
  
      for handle in handles:
          handle.remove()
      return result
  
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 5 for calling a global function with example usage: shape_inference(module) and returns: result