class_docstr: Execute an FX graph Node-by-Node with concrete tensor and record the memory
  usage, execution time of forward and backward, and type of the result into
  the corresponding node.
  
  Usage:
      BATCH_SIZE = 2
      DIM_IN = 4
      DIM_HIDDEN = 16
      DIM_OUT = 16
      model = torch.nn.Sequential(
          torch.nn.Linear(DIM_IN, DIM_HIDDEN),
          torch.nn.Linear(DIM_HIDDEN, DIM_OUT),
          ).cuda()
      input_sample = torch.rand(BATCH_SIZE, DIM_IN, device="cuda")
      gm = symbolic_trace(model)
      interp = ConcreteInfoProp(gm)
      interp.run(input_sample)
      print(interp.summary(unit='kb'))
  
  
      output of above code is
      Op type       Op             Forward time             Backward time    SAVE_FWD_IN    FWD_OUT    FWD_TMP    BWD_OUT    BWD_TMP
      -----------  -------  -----------------------  ------------------------  -------------  ---------  ---------  ---------  ---------
      placeholder  input_1                    0.0 s                     0.0 s          False    0.00 KB    0.00 KB    0.00 KB    0.00 KB
      call_module       _0  0.0003993511199951172 s     0.00706791877746582 s          False    0.50 KB    0.00 KB    0.03 KB    0.66 KB
      call_module       _1   6.29425048828125e-05 s  0.00018286705017089844 s          False    0.50 KB    0.00 KB    0.12 KB    0.81 KB
           output   output                    0.0 s                     0.0 s           True    0.00 KB    0.00 KB    0.00 KB    0.00 KB
  Args:
       module (GraphModule): The module to be executed
class_docstr_tok: ["Execute","an","FX","graph","Node-by-Node","with","concrete","tensor","and","record","the","memory","usage",",","execution","time","of","forward","and","backward",",","and","type","of","the","result","into","the","corresponding","node",".","Usage",":","BATCH_SIZE","=","2","DIM_IN","=","4","DIM_HIDDEN","=","16","DIM_OUT","=","16","model","=","torch.nn.Sequential","(","torch.nn.Linear","(","DIM_IN",",","DIM_HIDDEN",")",",","torch.nn.Linear","(","DIM_HIDDEN",",","DIM_OUT",")",",",")",".cuda","(",")","input_sample","=","torch.rand","(","BATCH_SIZE",",","DIM_IN",",","device=","''","cuda","''",")","gm","=","symbolic_trace","(","model",")","interp","=","ConcreteInfoProp","(","gm",")","interp.run","(","input_sample",")","print","(","interp.summary","(","unit='kb","'",")",")","output","of","above","code","is","Op","type","Op","Forward","time","Backward","time","SAVE_FWD_IN","FWD_OUT","FWD_TMP","BWD_OUT","BWD_TMP","--","--","--","--","--","-","--","--","--","-","--","--","--","--","--","--","--","--","--","--","--","-","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","-","--","--","--","--","-","--","--","--","--","-","--","--","--","--","-","--","--","--","--","-","placeholder","input_1","0.0","s","0.0","s","False","0.00","KB","0.00","KB","0.00","KB","0.00","KB","call_module","_0","0.0003993511199951172","s","0.00706791877746582","s","False","0.50","KB","0.00","KB","0.03","KB","0.66","KB","call_module","_1","6.29425048828125e-05","s","0.00018286705017089844","s","False","0.50","KB","0.00","KB","0.12","KB","0.81","KB","output","output","0.0","s","0.0","s","True","0.00","KB","0.00","KB","0.00","KB","0.00","KB","Args",":","module","(","GraphModule",")",":","The","module","to","be","executed"]
docstr: Summarizes the memory and FLOPs statistics of the `GraphModule` in
  tabular format. Note that this API requires the ``tabulate`` module
  to be installed.
docstr_tok: ["Summarizes","the","memory","and","FLOPs","statistics","of","the","`","GraphModule","`","in","tabular","format",".","Note","that","this","API","requires","the","``","tabulate","``","module","to","be","installed","."]
code: def summary(self, unit: str = "MB") -> str:
      """
      Summarizes the memory and FLOPs statistics of the `GraphModule` in
      tabular format. Note that this API requires the ``tabulate`` module
      to be installed.
      """
      # https://github.com/pytorch/pytorch/blob/master/torch/fx/graph.py
      try:
          from tabulate import tabulate
      except ImportError:
          print(
              "`summary` relies on the library `tabulate`, "
              "which could not be found on this machine. Run `pip "
              "install tabulate` to install the library."
          )
  
      assert (
          self._is_proped
      ), "Please call `interp.run(input)` before calling `interp.summary()`."
  
      # Build up a list of summary information for each node
      node_summaries: List[List[Any]] = []
  
      def mem_repr(mem: int) -> str:
          unit_divisor_map = {
              "kb": 1024,
              "mb": 1024**2,
              "gb": 1024**3,
              "tb": 1024**4,
          }
          return f"{mem / unit_divisor_map[unit.lower()]:.2f} {unit.upper()}"
  
      def time_repr(time: float):
          return f"{time:,} s"
  
      for node in self.module.graph.nodes:
          node: Node
          node_summaries.append(
              [
                  node.op,
                  str(node),
                  time_repr(node.meta["fwd_time"]),
                  time_repr(node.meta["bwd_time"]),
                  node.meta["save_fwd_in"],
                  mem_repr(node.meta["fwd_mem_out"]),
                  mem_repr(node.meta["fwd_mem_tmp"]),
                  mem_repr(node.meta["bwd_mem_out"]),
                  mem_repr(node.meta["bwd_mem_tmp"]),
              ]
          )
  
      # Use the ``tabulate`` library to create a well-formatted table
      # presenting our summary information
      headers: List[str] = [
          "Op type",
          "Op",
          "Forward time",
          "Backward time",
          "SAVE_FWD_IN",
          "FWD_OUT",
          "FWD_TMP",
          "BWD_OUT",
          "BWD_TMP",
      ]
  
      return tabulate(node_summaries, headers=headers, stralign="right")
  
file_path: colossalai/colossalai/fx/passes/concrete_info_prop.py
imports: from dataclasses import asdict
  from typing import Any, Dict, List, Optional, Tuple
  import torch
  import torch.fx
  from torch.fx.node import Argument, Node, Target
  from torch.utils._pytree import tree_flatten
  from colossalai.fx._compatibility import compatibility
  from colossalai.fx.profiler import GraphInfo, profile_function, profile_method, profile_module
label_desc: Use image node_id 10 for calling the ConcreteInfoProp obj's underlying member method code with example usage: obj.summary(unit) and returns: tabulate, unknown, str, unit, time, str