class_docstr: Estimate memory with chunk
class_docstr_tok: ["Estimate","memory","with","chunk"]
docstr: Estimate inference memory with chunk
  
  Args:
      node_list (List): _description_
      chunk_infos (Dict): Chunk information. Defaults to None.
      print_mem (bool): Wether to print peak memory of every node. Defaults to False.
  
  Returns:
      act_memory_peak_log (List): peak memory of every node
      act_memory_after_node_log (List): memory after executing every node
      active_node_list_log (List): active nodes of every node. active nodes refer to
          nodes generated but not deleted.
docstr_tok: ["Estimate","inference","memory","with","chunk","Args",":","node_list","(","List",")",":","_description_","chunk_infos","(","Dict",")",":","Chunk","information",".","Defaults","to","None",".","print_mem","(","bool",")",":","Wether","to","print","peak","memory","of","every","node",".","Defaults","to","False",".","Returns",":","act_memory_peak_log","(","List",")",":","peak","memory","of","every","node","act_memory_after_node_log","(","List",")",":","memory","after","executing","every","node","active_node_list_log","(","List",")",":","active","nodes","of","every","node",".","active","nodes","refer","to","nodes","generated","but","not","deleted","."]
code: def estimate_chunk_inference_mem(
      self,
      node_list: List,
      chunk_infos: Dict = None,
      print_mem: bool = False,
  ):
      """
      Estimate inference memory with chunk
  
      Args:
          node_list (List): _description_
          chunk_infos (Dict): Chunk information. Defaults to None.
          print_mem (bool): Wether to print peak memory of every node. Defaults to False.
  
      Returns:
          act_memory_peak_log (List): peak memory of every node
          act_memory_after_node_log (List): memory after executing every node
          active_node_list_log (List): active nodes of every node. active nodes refer to
              nodes generated but not deleted.
      """
      act_memory = 0.0
      act_memory_peak_log = []
      act_memory_after_node_log = []
      active_nodes = {}
      active_nodes_log = []
      not_contiguous_list = []
      node_mgr = NodeMgr(node_list)
      delete_node_dict = self._build_delete_node_dict(node_mgr)
  
      use_chunk = True if chunk_infos is not None else False
      chunk_within = False
      chunk_region_idx = None
      chunk_ratio = 1  # use it to estimate chunk mem
      chunk_inputs_all = []
  
      if use_chunk:
          chunk_regions = [i["region"] for i in chunk_infos]
          chunk_starts = [i[0] for i in chunk_regions]
          chunk_ends = [i[1] for i in chunk_regions]
          chunk_inputs = [i["inputs"] for i in chunk_infos]
          chunk_inputs_non_chunk = [
              i["inputs_non_chunk"] for i in chunk_infos
          ]
          chunk_inputs_all = [j for i in chunk_inputs for j in i] + [
              j for i in chunk_inputs_non_chunk for j in i
          ]
          chunk_outputs = [i["outputs"] for i in chunk_infos]
          chunk_node_dim = [i["node_chunk_dim"] for i in chunk_infos]
          chunk_sizes = [
              i["chunk_size"] if "chunk_size" in i else 1
              for i in chunk_infos
          ]
  
      for idx, node in enumerate(node_mgr.get_node_list()):
          # if node in chunk start nodes, change chunk ratio and add chunk_tensor
          if use_chunk and idx in chunk_starts:
              chunk_within = True
              chunk_region_idx = chunk_starts.index(idx)
              self._add_active_nodes_from_list(
                  active_nodes, chunk_outputs[chunk_region_idx]
              )
  
          # determine chunk ratio for current node
          if chunk_within:
              chunk_ratio = self._get_chunk_ratio(
                  node,
                  chunk_node_dim[chunk_region_idx],
                  chunk_sizes[chunk_region_idx],
              )
  
          # add current node as active node
          self._add_active_node(node, active_nodes, chunk_ratio)
          act_memory = self._get_memory_from_active_nodes(active_nodes)
  
          # if node is placeholder, just add the size of the node
          if node.op == "placeholder":
              act_memory_peak_log.append(act_memory)
          # skip output
          elif node.op == "output":
              continue
          # no change for non compute node
          elif is_non_memory_node(node):
              act_memory_peak_log.append(act_memory)
          # node is a compute op, calculate tmp
          else:
              # forward memory
              # TODO: contiguous_memory still not accurate for matmul, view, reshape and transpose
              tmp_memory = (
                  self._get_tmp_memory(
                      node, not_contiguous_list, delete=True
                  )
                  * chunk_ratio
              )
              # record max act memory
              act_memory_peak_log.append(act_memory + tmp_memory)
  
          # remove_deactive_node
          self._remove_deactive_node(
              idx,
              node,
              active_nodes,
              delete_node_dict,
              kept_nodes=chunk_inputs_all,
          )
  
          # if node in chunk end nodes, restore chunk settings
          if use_chunk and idx in chunk_ends:
              self._remove_deactive_node(
                  idx, node, active_nodes, delete_node_dict
              )  # dont provide kept nodes now
              chunk_within = False
              chunk_ratio = 1
              chunk_region_idx = None
  
          act_memory = self._get_memory_from_active_nodes(active_nodes)
          act_memory_after_node_log.append(act_memory)
          active_nodes_log.append(active_nodes.copy())
  
      if print_mem:
          print("with chunk" if use_chunk else "without chunk")
          self._print_compute_op_mem_log(
              act_memory_peak_log, node_mgr.get_node_list(), "peak"
          )
  
      # param_memory = parameter_size(gm)
      # all_memory = act_memory + param_memory
      return (
          act_memory_peak_log,
          act_memory_after_node_log,
          active_nodes_log,
      )
  
file_path: ColossalAI/colossalai/autochunk/estimate_memory.py
imports: from typing import Dict, List
  import torch
  from torch.fx.node import Node
  from .utils import NodeMgr, get_node_shape, is_non_memory_node
label_desc: Use image node_id 11 for calling the EstimateMemory obj's underlying member method code with example usage: obj.estimate_chunk_inference_mem(node_list, chunk_infos, print_mem) and returns: act_memory_peak_log,  act_memory_after_node_log,  active_nodes_log