code: def register_shape_inference_formula(
      class_or_func: Any, formula: Formula
  ) -> None:
      """
      Register a shape inference formula for a module.
  
      Parameters
      ----------
      class_or_func
          The module or function to register the formula for.
          The class here needs to be a class, not an instantiated module.
      formula
          A function that takes in a module and its inputs, and returns the output shape.
          To be specific, its input will be the module or function itself,
          plus ``*args`` and ``**kwargs`` of the module or function.
          Tensors will be replaced with :class:`ShapeTensor` objects.
          The output should be the same format as its normal forward output,
          but every tensors should be replaced with a :class:`MutableShape` object.
  
      Examples
      --------
      Here is an example of a formula for FC::
  
          def linear_formula(module: nn.Linear, input: ShapeTensor) -> MutableShape:
              return MutableShape(*tuple(input.real_shape)[:-1], module.out_features)
  
      It can be registered with::
  
          register_shape_inference_formula(nn.Linear, linear_formula)
      """
      if class_or_func in _shape_inference_formulas:
          _logger.warning(
              f"Overwriting shape inference formula for {class_or_func}"
          )
      _shape_inference_formulas[class_or_func] = formula
  
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 1 for calling a global function with example usage: register_shape_inference_formula(class_or_func, formula) without return types