code: def find_shape_inference_formula(
      module_or_func: Any,
  ) -> Formula | None:
      """
      Find the shape inference formula for a module or function.
  
      It searches two places in order:
  
      1. The module's ``_shape_forward`` attribute.
         The function should follow the signature defined in :func:`register_shape_inference_formula`.
      2. The global registry. Register with :func:`register_shape_inference_formula`.
  
      Parameters
      ----------
      module_or_func
          The module or function to find the formula for.
          The module here needs to be an instantiated module, not a class.
      """
  
      if isinstance(module_or_func, nn.Module):
          formula = None
          # Shape inference formulas are fetched in the following order:
          # 1. Check `_shape_forward` of the module.
          if hasattr(module_or_func.__class__, "_shape_forward"):
              formula: Any = module_or_func.__class__._shape_forward  # type: ignore
          # 2. Check global registered formulas.
          elif type(module_or_func) in _shape_inference_formulas:
              formula = _shape_inference_formulas[type(module_or_func)]
          return formula
      else:
          return _shape_inference_formulas.get(module_or_func)
  
imports: from __future__ import annotations
  import logging
  import warnings
  from typing import Callable, Type, Tuple, Any, cast
  import torch
  from torch import nn
  import nni.nas.nn.pytorch
  from nni.mutable import MutableExpression
  from .shape import Formula, ShapeTensor, MutableShape, extract_shape_info, switch_case_shape_info, shape_inference
  from ._attrs import _getattr, tuple_2_t
label_desc: Use image node_id 2 for calling a global function with example usage: find_shape_inference_formula(module_or_func) and returns: formula, _shape_inference_formulas