class_docstr: A context manager to measure flops of a module. Requires PyTorch 1.13+.
  
  Flop count implementation based on
  https://dev-discuss.pytorch.org/t/the-ideal-pytorch-flop-counter-with-torch-dispatch/505
  
  Examples::
  
      >>> import copy
      >>> import torch
      >>> import torchvision.models as models
      >>> from torchtnt.utils.flops import FlopTensorDispatchMode
  
      >>> module = models.resnet18()
      >>> module_input = torch.randn(1, 3, 224, 224)
      >>> with FlopTensorDispatchMode(module) as ftdm:
      >>>     # count forward flops
      >>>     res = module(module_input).mean()
      >>>     flops_forward = copy.deepcopy(ftdm.flop_counts)
  
      >>>     # reset count before counting backward flops
      >>>     ftdm.reset()
      >>>     res.backward()
      >>>     flops_backward = copy.deepcopy(ftdm.flop_counts)
class_docstr_tok: ["A","context","manager","to","measure","flops","of","a","module",".","Requires","PyTorch","1.13+",".","Flop","count","implementation","based","on","https",":","\/\/dev-discuss.pytorch.org\/t\/the-ideal-pytorch-flop-counter-with-torch-dispatch\/505","Examples",":",":",">",">",">","import","copy",">",">",">","import","torch",">",">",">","import","torchvision.models","as","models",">",">",">","from","torchtnt.utils.flops","import","FlopTensorDispatchMode",">",">",">","module","=","models.resnet18","(",")",">",">",">","module_input","=","torch.randn","(","1",",","3",",","224",",","224",")",">",">",">","with","FlopTensorDispatchMode","(","module",")","as","ftdm",":",">",">",">","#","count","forward","flops",">",">",">","res","=","module","(","module_input",")",".mean","(",")",">",">",">","flops_forward","=","copy.deepcopy","(","ftdm.flop_counts",")",">",">",">","#","reset","count","before","counting","backward","flops",">",">",">","ftdm.reset","(",")",">",">",">","res.backward","(",")",">",">",">","flops_backward","=","copy.deepcopy","(","ftdm.flop_counts",")"]
code: def __torch_dispatch__(
      self,
      func: Callable[..., Any],  # pyre-ignore [2] func can be any func
      types: Tuple[Any],  # pyre-ignore [2]
      args=(),  # pyre-ignore [2]
      kwargs=None,  # pyre-ignore [2]
  ) -> PyTree:
      rs = func(*args, **kwargs)
      outs = _normalize_tuple(rs)
  
      if func in flop_mapping:
          flop_count = flop_mapping[func](args, outs)
          for par in self._parents:
              # pyre-ignore [58]
              self.flop_counts[par][func.__name__] += flop_count
      else:
          logging.debug(
              f"{func} is not yet supported in FLOPs calculation."
          )
  
      return rs
  
imports: import logging
  import operator
  from collections import defaultdict
  from functools import reduce
  from numbers import Number
  from typing import Any, Callable, DefaultDict, Dict, List, Tuple
  import torch
  from torch.utils._python_dispatch import TorchDispatchMode
  from torch.utils._pytree import PyTree, tree_map
label_desc: Use image node_id 3 for calling the FlopTensorDispatchMode obj's underlying member method code with example usage: obj.__torch_dispatch__(func, types, args, kwargs) and returns: rs