code: def _get_all_node_info(self, end_dim, start_idx, end_idx):
      cur_node_list = [
          self.node_mgr.get_node_by_idx(end_idx)
      ]  # start from the last node
      all_node_info = {
          cur_node_list[0]: {"chunk_dim": end_dim, "fix_dim": []}
      }
  
      while len(cur_node_list) > 0:
          next_node_list = []
  
          for cur_node in cur_node_list:
              # get cur node info
              cur_node_chunk_dim = all_node_info[cur_node]["chunk_dim"]
              cur_node_fix_dim = all_node_info[cur_node]["fix_dim"]
              if cur_node_chunk_dim is not None:
                  cur_node_compute = (
                      self.trace_indice._find_compute_trace_from_node(
                          cur_node
                      )
                  )
                  cur_node_source = (
                      self.trace_indice._find_source_trace_from_node(
                          cur_node
                      )
                  )
              else:
                  cur_node_compute = cur_node_source = None
  
              # get all valid args
              arg_list = []
              for arg in cur_node.all_input_nodes:
                  if type(arg) != type(cur_node):
                      continue
                  if is_non_compute_node(arg):
                      continue
                  if get_node_shape(arg) is None:
                      continue
                  arg_list.append(arg)
                  flow_flag = self._assign_single_node_flow(
                      arg,
                      start_idx,
                      end_idx,
                      cur_node,
                      cur_node_chunk_dim,
                      cur_node_compute,
                      cur_node_source,
                      cur_node_fix_dim,
                      all_node_info,
                      next_node_list,
                  )
                  if flow_flag == False:
                      return None
  
          cur_node_list = next_node_list
      return all_node_info
  
file_path: colossalai/colossalai/autochunk/trace_flow.py
imports: from typing import Dict, List, Tuple
  from torch.fx.node import Node
  from .trace_indice import TraceIndice
  from .utils import NodeMgr, find_chunk_all_input_nodes, find_chunk_compute_input_and_output_nodes, find_tensor_shape_node, flat_list, get_node_name, get_node_shape, is_non_compute_node
label_desc: Use image node_id 5 for calling the TraceFlow obj's underlying member method code with example usage: obj._get_all_node_info(end_dim, start_idx, end_idx) and returns: all_node_info, None