class_docstr: A tensor with mutable shape information.
  Inspired by `Ideal FLOPS counter <https://dev-discuss.pytorch.org/t/the-ideal-pytorch-flop-counter-with-torch-dispatch/505>`__.
class_docstr_tok: ["A","tensor","with","mutable","shape","information",".","Inspired","by","`","Ideal","FLOPS","counter","<","https",":","\/\/dev-discuss.pytorch.org\/t\/the-ideal-pytorch-flop-counter-with-torch-dispatch\/505",">","`","__","."]
code: def __torch_dispatch__(cls, func, types, args=(), kwargs=None):
      from torch.utils._pytree import tree_map
      from .shape_formula import find_shape_inference_formula
  
      kwargs = cast(Dict[str, Any], kwargs)
  
      def unwrap(e):
          return e.elem if isinstance(e, ShapeTensor) else e
  
      def wrap(e):
          return (
              ShapeTensor(e)
              if isinstance(e, torch.Tensor)
              and not isinstance(e, ShapeTensor)
              else e
          )
  
      rs = func(*tree_map(unwrap, args), **tree_map(unwrap, kwargs))
      rs = tree_map(wrap, rs)
  
      if (
          _current_module_names
          and _current_module_names[-1].has_formula
      ):
          # The closest module has a shape formula, so we don't need to infer the shape.
          return rs
  
      formula = find_shape_inference_formula(func)
  
      def format_module_stack():
          current_modules_info = [
              f"  - {m.name!r} (type: {m.type_name}, "
              + ("NO" if not m.has_formula else "HAS")
              + " shape formula)"
              for m in _current_module_names
          ]
          return "\n".join(current_modules_info)
  
      if formula is not None:
          try:
              shapes = formula(func, *args, **kwargs)
              rs = _assign_shape_info(rs, shapes)
          except:
              _logger.warning(
                  "Shape inference of %s failed though shape inference formula found: %s. Module calling stack:\n%s\n"
                  "This is either a bug in the shape inference formula, or something goes wrong earlier. "
                  "To workaround this, you can explicitly write a formula for its parent modules.\n%s",
                  func,
                  formula,
                  format_module_stack(),
                  traceback.format_exc(),
              )
      else:
          if not any(m.has_formula for m in _current_module_names):
              msg = "and none of its parent modules have a shape inference formula. Shape inference is likely to fail. "
          else:
              msg = "and a recent module that needs shape information has no shape inference formula. "
          _logger.warning(
              f"Shape information is not explicitly propagated when executing {func}, and {msg}"
              "Module calling stack:\n" + format_module_stack()
          )
  
      return rs
  
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 3 for calling the ShapeTensor obj's underlying member method code with example usage: obj.__torch_dispatch__(cls, func, types, args, kwargs) and returns: rs, unknown, unknown, rs, str