class_docstr: Class to represent node in Markov decision process. A stage has states
  to represent different schedules of the current node. Since in this problem
  the action is the schedule selected for current node, action can be fully
  represented by states. No extra attribute needs for action.
  
  In most cases, instance of this class should be created through DPTuner.
class_docstr_tok: ["Class","to","represent","node","in","Markov","decision","process",".","A","stage","has","states","to","represent","different","schedules","of","the","current","node",".","Since","in","this","problem","the","action","is","the","schedule","selected","for","current","node",",","action","can","be","fully","represented","by","states",".","No","extra","attribute","needs","for","action",".","In","most","cases",",","instance","of","this","class","should","be","created","through","DPTuner","."]
docstr: Align all input node states shapes to be the same and transpose/reshape properly.
  
  This is used in creating multi_input operator states.
  
  Parameters
  ----------
  input_index_list : list of int
      List of input node index.
  
  stage_dict : dict of int to Stage
      Global dictionary of node index to stage.
  
  node_list : list of dict
      List of all nodes for current graph.
  
  Returns
  -------
  states_list : list of tuple
      List of aligned states.
  
  aligned_node_list : list in int
      List of node index for aligned states.
docstr_tok: ["Align","all","input","node","states","shapes","to","be","the","same","and","transpose\/reshape","properly",".","This","is","used","in","creating","multi_input","operator","states",".","Parameters","--","--","--","--","--","input_index_list",":","list","of","int","List","of","input","node","index",".","stage_dict",":","dict","of","int","to","Stage","Global","dictionary","of","node","index","to","stage",".","node_list",":","list","of","dict","List","of","all","nodes","for","current","graph",".","Returns","--","--","--","-","states_list",":","list","of","tuple","List","of","aligned","states",".","aligned_node_list",":","list","in","int","List","of","node","index","for","aligned","states","."]
code: def align_states(input_index_list, stage_dict, node_list):
      """Align all input node states shapes to be the same and transpose/reshape properly.
  
      This is used in creating multi_input operator states.
  
      Parameters
      ----------
      input_index_list : list of int
          List of input node index.
  
      stage_dict : dict of int to Stage
          Global dictionary of node index to stage.
  
      node_list : list of dict
          List of all nodes for current graph.
  
      Returns
      -------
      states_list : list of tuple
          List of aligned states.
  
      aligned_node_list : list in int
          List of node index for aligned states.
      """
      aligned_node_list = list(input_index_list)
      states_list = []
      for input_idx in input_index_list:
          input_node_stage = stage_dict[input_idx]
          for dep_idx in input_node_stage.dep:
              if dep_idx not in aligned_node_list:
                  aligned_node_list.append(dep_idx)
      aligned_shape = []
      for idx in aligned_node_list:
          aligned_shape.append(len(node_list[idx]["record_candidates"]))
      for input_idx in input_index_list:
          input_node_stage = stage_dict[input_idx]
          input_node_shape_idx_list = [input_idx] + input_node_stage.dep
          transpose_idx_list = []
          reshape_list = []
          major_axis = -1
          for i, idx in enumerate(aligned_node_list):
              if input_idx == idx:
                  major_axis = i
              if idx in input_node_shape_idx_list:
                  transpose_idx_list.append(idx)
                  reshape_list.append(aligned_shape[i])
              else:
                  reshape_list.append(1)
          transpose_list = [
              input_node_shape_idx_list.index(idx)
              for idx in transpose_idx_list
          ]
          input_node_states = np.transpose(
              input_node_stage.states, tuple(transpose_list)
          )
          input_node_states = np.reshape(
              input_node_states, tuple(reshape_list)
          )
          input_node_states = np.broadcast_to(
              input_node_states, aligned_shape
          )
          multiplier = 1
          for i in range(major_axis + 1, len(aligned_shape)):
              multiplier *= aligned_shape[i]
          states_list.append(
              (input_idx, major_axis, multiplier, input_node_states)
          )
      return states_list, aligned_node_list
  
file_path: tvm/python/tvm/autotvm/graph_tuner/dynamic_programming_stage.py
imports: import numpy
  from .utils import is_boundary_node
label_desc: Use image node_id 9 for calling the DPStage obj's underlying member method code with example usage: obj.align_states(input_index_list, stage_dict, node_list) and returns: states_list,  aligned_node_list