code: def _module_shape_inference_impl(
      module: nn.Module,
      outputs: Any,
      *input_args: ShapeTensor,
      is_leaf: Callable[[nn.Module], bool] | None = None,
      **input_kwargs: ShapeTensor,
  ) -> ShapeTensor | tuple[ShapeTensor, ...]:
      from torch.utils._pytree import tree_map
  
      from .shape_formula import find_shape_inference_formula
  
      formula = find_shape_inference_formula(module)
  
      if formula is None:
          # No formula, but the shapes could've been already inferred.
          def _ensure_shape(tensor: ShapeTensor) -> None:
              if (
                  isinstance(tensor, ShapeTensor)
                  and tensor.real_shape is None
              ):
                  # TODO: more informative error message
                  module_repr = repr(module)
                  if len(module_repr) > 100:
                      module_repr = module_repr[:100] + "..."
                  raise RuntimeError(
                      f"Shape inference failed because no shape inference formula is found for {module_repr} "
                      f"of type {type(module).__name__}. "
                      "Meanwhile the nested modules and functions inside failed to propagate the shape information. "
                      "Please provide a `_shape_forward` member function or register a formula using `register_shape_inference_formula`."
                  )
  
          # Check every ShapeTensor in outputs
          tree_map(_ensure_shape, outputs)
          return outputs
  
      # If we find a formula, we must run the formula to find out the shape,
      # because the shapes attached to the tensor could not be desired.
  
      # Some formulas support an optional argument `is_leaf`.
      formula_kwargs = {}
      if argument_in_spec(formula, "is_leaf"):
          formula_kwargs["is_leaf"] = is_leaf
  
      output_shape = formula(
          module, *input_args, **formula_kwargs, **input_kwargs
      )
      outputs = _assign_shape_info(outputs, output_shape)
      return outputs
  
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 8 for calling a global function with example usage: _module_shape_inference_impl(module, outputs) and returns: outputs, outputs