class_docstr: This is the core class for AutoChunk.
  
  It defines the framework of the strategy of AutoChunk.
  Chunks will be selected one by one until search stops.
  
  The chunk search is as follows:
  1. find the peak memory node
  2. find the max chunk region according to the peak memory node
  3. find all possible chunk regions in the max chunk region
  4. find the best chunk region for current status
  5. goto 1
  
  Attributes:
      gm: graph model
      print_mem (bool): print estimated memory
      trace_index: trace the flow of every dim of every node to find all free dims
      trace_flow: determine the region chunk strategy
      reorder_graph: reorder nodes to improve chunk efficiency
      estimate_memory: estimate memory with chunk
      select_chunk: select the best chunk region
  
  Args:
      gm: graph model
      max_memory (int): max memory in MB
      print_mem (bool): print estimated memory
class_docstr_tok: ["This","is","the","core","class","for","AutoChunk",".","It","defines","the","framework","of","the","strategy","of","AutoChunk",".","Chunks","will","be","selected","one","by","one","until","search","stops",".","The","chunk","search","is","as","follows",":","1.","find","the","peak","memory","node","2.","find","the","max","chunk","region","according","to","the","peak","memory","node","3.","find","all","possible","chunk","regions","in","the","max","chunk","region","4.","find","the","best","chunk","region","for","current","status","5.","goto","1","Attributes",":","gm",":","graph","model","print_mem","(","bool",")",":","print","estimated","memory","trace_index",":","trace","the","flow","of","every","dim","of","every","node","to","find","all","free","dims","trace_flow",":","determine","the","region","chunk","strategy","reorder_graph",":","reorder","nodes","to","improve","chunk","efficiency","estimate_memory",":","estimate","memory","with","chunk","select_chunk",":","select","the","best","chunk","region","Args",":","gm",":","graph","model","max_memory","(","int",")",":","max","memory","in","MB","print_mem","(","bool",")",":","print","estimated","memory"]
docstr: Find chunk info for a region.
  
  We are given the region start and region end, and need to find out all chunk info for it.
  We first loop every dim of start node and end node, to see if we can find dim pair,
  which is linked in a flow and not computed.
  If found, we then search flow in the whole region to find out all chunk infos.
  
  Args:
      input_trace (List): node's input trace in region
      output_trace (List): node's output trace in region
      start_idx (int): region start node index
      end_idx (int): region end node index
  
  Returns:
      chunk_infos: possible regions found
docstr_tok: ["Find","chunk","info","for","a","region",".","We","are","given","the","region","start","and","region","end",",","and","need","to","find","out","all","chunk","info","for","it",".","We","first","loop","every","dim","of","start","node","and","end","node",",","to","see","if","we","can","find","dim","pair",",","which","is","linked","in","a","flow","and","not","computed",".","If","found",",","we","then","search","flow","in","the","whole","region","to","find","out","all","chunk","infos",".","Args",":","input_trace","(","List",")",":","node","'s","input","trace","in","region","output_trace","(","List",")",":","node","'s","output","trace","in","region","start_idx","(","int",")",":","region","start","node","index","end_idx","(","int",")",":","region","end","node","index","Returns",":","chunk_infos",":","possible","regions","found"]
code: def _find_chunk_info(
      self, input_trace, output_trace, start_idx, end_idx
  ) -> List:
      """
      Find chunk info for a region.
  
      We are given the region start and region end, and need to find out all chunk info for it.
      We first loop every dim of start node and end node, to see if we can find dim pair,
      which is linked in a flow and not computed.
      If found, we then search flow in the whole region to find out all chunk infos.
  
      Args:
          input_trace (List): node's input trace in region
          output_trace (List): node's output trace in region
          start_idx (int): region start node index
          end_idx (int): region end node index
  
      Returns:
          chunk_infos: possible regions found
      """
      start_traces = input_trace[start_idx]
      if len(start_traces) > 1:  # TODO need to be removed
          return []
      end_trace = output_trace[end_idx]
      end_node = self.node_mgr.get_node_by_idx(end_idx)
  
      chunk_infos = []
      for end_dim, _ in enumerate(end_trace["indice"]):
          for start_node, start_trace in start_traces.items():
              for start_dim, _ in enumerate(start_trace["indice"]):
                  if not self.trace_flow.check_region_start_end(
                      start_node,
                      start_dim,
                      start_idx,
                      end_node,
                      end_dim,
                      end_idx,
                  ):
                      continue
                  # flow search
                  chunk_info = self.trace_flow.flow_search(
                      start_idx, start_dim, end_idx, end_dim
                  )
                  if chunk_info is None:
                      continue
                  chunk_infos.append(chunk_info)
      return chunk_infos
  
file_path: colossalai/colossalai/autochunk/search_chunk.py
imports: import copy
  from typing import Dict, List, Tuple
  from torch.fx.node import Node
  from .estimate_memory import EstimateMemory
  from .reorder_graph import ReorderGraph
  from .select_chunk import SelectChunk
  from .trace_flow import TraceFlow
  from .trace_indice import TraceIndice
  from .utils import NodeMgr, get_logger, is_non_compute_node, is_non_compute_node_except_placeholder
label_desc: Use image node_id 5 for calling the SearchChunk obj's underlying member method code with example usage: obj._find_chunk_info(input_trace, output_trace, start_idx, end_idx) and returns: chunk_infos, list