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 _create_backwards_pop(self, name: str) -> Callable[..., Any]:
      class PopState(torch.autograd.Function):
          @staticmethod
          def forward(ctx, *args):
              args = tree_map(
                  lambda x: x.clone()
                  if isinstance(x, torch.Tensor)
                  else x,
                  args,
              )
              if len(args) == 1:
                  return args[0]
              return args
  
          @staticmethod
          def backward(ctx, *grad_outs):
              parents = self._parents
              assert parents[-1] == name
              parents.pop()
              return grad_outs
  
      # Pyre does not support analyzing classes nested in functions.
      # But this class can't be lifted out of the function as it is a static class
      # using a function parameter.
      return PopState.apply
  
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 5 for calling the FlopTensorDispatchMode obj's underlying member method code with example usage: obj._create_backwards_pop(name) and returns: PopState, args, grad_outs, args