code: def flop_count(
      module: Union[torch.nn.Module, Callable],
      *args,
      verbose: bool = False,
      forward_only: bool = True,
      **kwargs,
  ) -> int:
      """
      Count the number of floating point operations in a model.
      Ideas from https://pastebin.com/AkvAyJBw.
      Parameters
      ----------
      module : Union[torch.nn.Module, Callable]
          The model to count the number of floating point operations.
      args : Any
          The positional arguments to pass to the model.
      verbose : bool
          Whether to print the number of floating point operations.
      forward_only : bool
          Whether to only count the number of floating point operations in the forward pass.
      kwargs : Any
          The keyword arguments to pass to the model.
  
      Returns
      -------
      int
          The number of floating point operations.
      """
      maybe_inplace = (
          getattr(module, "inplace", False)
          or kwargs.get("inplace", False)
          or getattr(module, "__name__", None)
          in ("add_", "mul_", "div_", "sub_")
      )
  
      class DummyModule(torch.nn.Module):
          def __init__(self, func):
              super().__init__()
              self.func = func
              self.__name__ = func.__name__
  
          def forward(self, *args, **kwargs):
              return self.func(*args, **kwargs)
  
      total_flop_count = {Phase.FWD: 0, Phase.BWD: 0}
      flop_counts = defaultdict(lambda: defaultdict(int))
      parents = ["Global"]
      module = (
          module
          if isinstance(module, torch.nn.Module)
          else DummyModule(module)
      )
  
      class FlopTensor(torch.Tensor):
          elem: torch.Tensor
  
          __slots__ = ["elem"]
  
          @staticmethod
          def __new__(cls, elem):
              # The wrapping tensor (FlopTensor) shouldn't hold any
              # memory for the class in question, but it should still
              # advertise the same device as before
              r = torch.Tensor._make_wrapper_subclass(  # type: ignore[attr-defined]
                  cls,
                  elem.size(),
                  strides=elem.stride(),
                  storage_offset=elem.storage_offset(),
                  # TODO: clone storage aliasing
                  dtype=elem.dtype,
                  layout=elem.layout,
                  device=elem.device,
                  requires_grad=elem.requires_grad,
              )
              # ...the real tensor is held as an element on the tensor.
              r.elem = elem
              return r
  
          def __repr__(self):
              if self.grad_fn:
                  return (
                      f"FlopTensor({self.elem}, grad_fn={self.grad_fn})"
                  )
              return f"FlopTensor({self.elem})"
  
          @classmethod
          def __torch_dispatch__(
              cls, func, types, args=(), kwargs=None
          ):
              def unwrap(e):
                  return e.elem if isinstance(e, FlopTensor) else e
  
              # no_dispatch is only needed if you use enable_python_mode.
              # It prevents infinite recursion.
              rs = func(
                  *tree_map(unwrap, args), **tree_map(unwrap, kwargs)
              )
              outs = normalize_tuple(rs)
  
              if func in flop_mapping:
                  nonlocal flop_counts, total_flop_count, cur_phase
                  flop_count = flop_mapping.get(func, zero_flop_jit)(
                      args, outs
                  )
                  for par in parents:
                      flop_counts[par][func.__name__] += flop_count
                  total_flop_count[cur_phase] += flop_count
  
              def wrap(e):
                  return (
                      FlopTensor(e)
                      if isinstance(e, torch.Tensor)
                      else e
                  )
  
              rs = tree_map(wrap, rs)
              return rs
  
      def is_autogradable(x):
          return isinstance(x, torch.Tensor) and x.is_floating_point()
  
      def normalize_tuple(x):
          if not isinstance(x, tuple):
              return (x,)
          return x
  
      def create_backwards_push(name):
          class PushState(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):
                  nonlocal parents
                  parents.append(name)
                  return grad_outs
  
          return PushState.apply
  
      def create_backwards_pop(name):
          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):
                  nonlocal parents
                  assert parents[-1] == name
                  parents.pop()
                  return grad_outs
  
          return PopState.apply
  
      def enter_module(name):
          def f(module, inputs):
              nonlocal parents
              parents.append(name)
              inputs = normalize_tuple(inputs)
              out = create_backwards_pop(name)(*inputs)
              return out
  
          return f
  
      def exit_module(name):
          def f(module, inputs, outputs):
              nonlocal parents
              assert parents[-1] == name
              parents.pop()
              outputs = normalize_tuple(outputs)
              return create_backwards_push(name)(*outputs)
  
          return f
  
      @contextmanager
      def instrument_module(mod):
          registered = []
          for name, module in dict(mod.named_children()).items():
              registered.append(
                  module.register_forward_pre_hook(enter_module(name))
              )
              registered.append(
                  module.register_forward_hook(exit_module(name))
              )
          yield
          for handle in registered:
              handle.remove()
  
      def display_flops():
          for mod in flop_counts.keys():
              print(f"Module: ", mod)
              for k, v in flop_counts[mod].items():
                  print("\t", k, _format_flops(v))
              print()
  
      def detach_variables(r):
          if isinstance(r, torch.Tensor):
              requires_grad = r.requires_grad
              r = r.detach()
              r.requires_grad = requires_grad
          return r
  
      def wrap(r):
          if isinstance(r, torch.Tensor):
              r = FlopTensor(detach_variables(r))
              if maybe_inplace:
                  r = r + 0
          return r
  
      with instrument_module(module):
          cur_phase = Phase.FWD
          rst = module(*tree_map(wrap, args), **tree_map(wrap, kwargs))
          rst = tuple(
              r
              for r in normalize_tuple(rst)
              if is_autogradable(r) and r.requires_grad
          )
          cur_phase = Phase.BWD
  
          if rst and not forward_only:
              grad = [torch.zeros_like(t) for t in rst]
              torch.autograd.backward(
                  rst,
                  grad,
              )
  
      if verbose:
          display_flops()
  
      if forward_only:
          return total_flop_count[Phase.FWD]
      else:
          return (
              total_flop_count[Phase.FWD],
              total_flop_count[Phase.BWD],
          )
  
imports: import operator
  from collections import defaultdict
  from contextlib import contextmanager
  from enum import Enum, auto
  from functools import partial, reduce
  from typing import Any, Callable, List, Union
  import torch
  from torch.utils._pytree import tree_map
label_desc: Use image node_id 2 for calling a global function with example usage: flop_count(module) and returns: unknown, x, PushState, PopState, f, f, r, r, total_flop_count, total_flop_count,  total_flop_count, self, r, str, self, str, rs, x, out, create_backwards_push, str, self, str, self, str, unknown, unknown, args, grad_outs, args, grad_outs, args, args