class_docstr: ``FlopCounterMode`` is a context manager that counts the number of flops within its context.
  
  It does this using a ``TorchDispatchMode``.
  
  It also supports hierarchical output by passing a module (or list of
  modules) to FlopCounterMode on construction. If you do not need hierarchical
  output, you do not need to use it with a module.
  
  Example usage
  
  .. code-block:: python
  
      mod = ...
      flop_counter = FlopCounterMode(mod)
      with flop_counter:
          mod.sum().backward()
class_docstr_tok: ["``","FlopCounterMode","``","is","a","context","manager","that","counts","the","number","of","flops","within","its","context",".","It","does","this","using","a","``","TorchDispatchMode","``",".","It","also","supports","hierarchical","output","by","passing","a","module","(","or","list","of","modules",")","to","FlopCounterMode","on","construction",".","If","you","do","not","need","hierarchical","output",",","you","do","not","need","to","use","it","with","a","module",".","Example","usage","..","code-block",":",":","python","mod","=","...","flop_counter","=","FlopCounterMode","(","mod",")","with","flop_counter",":","mod.sum","(",")",".backward","(",")"]
code: def __init__(
      self,
      mods: Optional[
          Union[torch.nn.Module, List[torch.nn.Module]]
      ] = None,
      depth: int = 2,
      display: bool = True,
      custom_mapping: Optional[Dict[Any, Any]] = None,
  ):
      self.flop_counts: Dict[str, Dict[Any, int]] = defaultdict(
          lambda: defaultdict(int)
      )
      self.depth = depth
      self.parents = ["Global"]
      self.display = display
      if custom_mapping is None:
          custom_mapping = {}
      if isinstance(mods, torch.nn.Module):
          mods = [mods]
      self.mods = mods
      # Keys will include the modules in `mods` and their submodules
      self._module_to_forward_hook_handles: Dict[
          nn.Module, _ForwardHookHandles
      ] = {}
      self.flop_registry = {
          **flop_registry,
          **{
              k: v
              if getattr(v, "_get_raw", False)
              else shape_wrapper(v)
              for k, v in custom_mapping.items()
          },
      }
  
imports: import torch
  import torch.nn
  from torch.utils._pytree import tree_map, tree_flatten, tree_unflatten
  from typing import List, Any, Dict, Optional, Union, NamedTuple
  from collections import defaultdict
  from torch.utils._python_dispatch import TorchDispatchMode
  from torch.utils.hooks import RemovableHandle
  from torch._decomp import register_decomposition
  from math import prod
  from functools import wraps
label_desc: Use image node_id 1 to create a new FlopCounterMode object from inherited base classes: TorchDispatchMode with example: obj = FlopCounterMode(mods, depth, display, custom_mapping)