parameters: (self, serial_main_program)
docstring: Complete the annotation of vars and ops in the update phase for parallel program.
docstring_summary: Complete the annotation of vars and ops in the update phase for parallel program.
docstring_tokens: Complete the annotation of vars and ops in the update phase for parallel program .
function: def complete_update_annotation(self, serial_main_program):
          """Complete the annotation of vars and ops in the update phase for parallel program."""
          ops = list(serial_main_program.global_block().ops)
          vars = serial_main_program.global_block().vars
          learning_rate_completed = False
  
          for idx in range(len(ops)):
  
              # complete the annotation of the optimizer op.
              # TODO to add attribute for moment var
              op = ops[idx]
              if int(op.attr('op_role')) == int(OpRole.Optimize):
  
                  if "Grad" in op.input_names and "Param" in ops[idx].input_names:
                      assert len(op.input(
                          "Param")) == 1, "Only support one-to-one now."
                      assert len(op.input(
                          "Grad")) == 1, "Only support one-to-one now."
                      param = vars[op.input("Param")[0]]
                      grad_var = vars[op.input("Grad")[0]]
  
                      param_dist_attr = self._dist_context.get_tensor_dist_attr_for_program(
                          param)
                      assert param_dist_attr is not None
                      ref_process_mesh = self._dist_context.get_tensor_dist_attr_for_program(
                          param).process_mesh
                      assert ref_process_mesh is not None
                      ref_dims_mapping = self._dist_context.get_tensor_dist_attr_for_program(
                          param).dims_mapping
                      assert ref_dims_mapping is not None
                      op_dist_attr = OperatorDistributedAttribute()
                      op_dist_attr.process_mesh = ref_process_mesh
                      op_dist_attr.set_input_dims_mapping(grad_var.name,
                                                          ref_dims_mapping)
                      op_dist_attr.set_input_dims_mapping(param.name,
                                                          ref_dims_mapping)
                      op_dist_attr.set_output_dims_mapping(param.name,
                                                           ref_dims_mapping)
                      learning_var = vars[op.input("LearningRate")[0]]
                      op_dist_attr.set_input_dims_mapping(learning_var.name, [-1])
                      op_dist_attr.set_output_dims_mapping(learning_var.name,
                                                           [-1])
  
                      if not learning_rate_completed:
                          learning_rate_completed = True
                          var_dist_attr = TensorDistributedAttribute()
                          var_dist_attr.process_mesh = ref_process_mesh
                          var_dist_attr.dims_mapping = [-1]
                          self._dist_context.set_tensor_dist_attr_for_program(
                              learning_var, var_dist_attr)
  
                      for input_name in op.desc.input_names():
  
                          if input_name in [
                                  'Param', 'Grad', 'LearningRate', "SkipUpdate",
                                  "Beta1Tensor", "Beta2Tensor", "EpsilonTensor",
                                  "MasterParam"
                          ]:
                              continue
  
                          assert len(op.desc.input(input_name)) == 1
                          input_var = vars[op.desc.input(input_name)[0]]
                          input_var_attr = TensorDistributedAttribute()
  
                          if "Beta1Pow" in input_name or "Beta2Pow" in input_name:
                              input_var_attr.dims_mapping = [-1]
                              op_dist_attr.set_input_dims_mapping(input_var.name,
                                                                  [-1])
                              op_dist_attr.set_output_dims_mapping(input_var.name,
                                                                   [-1])
                          else:
                              assert "Moment" in input_name
                              input_var_attr.dims_mapping = ref_dims_mapping
                              op_dist_attr.set_input_dims_mapping(
                                  input_var.name, ref_dims_mapping)
                              op_dist_attr.set_output_dims_mapping(
                                  input_var.name, ref_dims_mapping)
  
                          input_var_attr.process_mesh = ref_process_mesh
                          self._dist_context.set_tensor_dist_attr_for_program(
                              input_var, input_var_attr)
  
                      self._dist_context.set_op_dist_attr_for_program(
                          op, op_dist_attr)
                      continue
function_tokens: def complete_update_annotation ( self , serial_main_program ) : ops = list ( serial_main_program . global_block ( ) . ops ) vars = serial_main_program . global_block ( ) . vars learning_rate_completed = False for idx in range ( len ( ops ) ) : # complete the annotation of the optimizer op. # TODO to add attribute for moment var op = ops [ idx ] if int ( op . attr ( 'op_role' ) ) == int ( OpRole . Optimize ) : if "Grad" in op . input_names and "Param" in ops [ idx ] . input_names : assert len ( op . input ( "Param" ) ) == 1 , "Only support one-to-one now." assert len ( op . input ( "Grad" ) ) == 1 , "Only support one-to-one now." param = vars [ op . input ( "Param" ) [ 0 ] ] grad_var = vars [ op . input ( "Grad" ) [ 0 ] ] param_dist_attr = self . _dist_context . get_tensor_dist_attr_for_program ( param ) assert param_dist_attr is not None ref_process_mesh = self . _dist_context . get_tensor_dist_attr_for_program ( param ) . process_mesh assert ref_process_mesh is not None ref_dims_mapping = self . _dist_context . get_tensor_dist_attr_for_program ( param ) . dims_mapping assert ref_dims_mapping is not None op_dist_attr = OperatorDistributedAttribute ( ) op_dist_attr . process_mesh = ref_process_mesh op_dist_attr . set_input_dims_mapping ( grad_var . name , ref_dims_mapping ) op_dist_attr . set_input_dims_mapping ( param . name , ref_dims_mapping ) op_dist_attr . set_output_dims_mapping ( param . name , ref_dims_mapping ) learning_var = vars [ op . input ( "LearningRate" ) [ 0 ] ] op_dist_attr . set_input_dims_mapping ( learning_var . name , [ - 1 ] ) op_dist_attr . set_output_dims_mapping ( learning_var . name , [ - 1 ] ) if not learning_rate_completed : learning_rate_completed = True var_dist_attr = TensorDistributedAttribute ( ) var_dist_attr . process_mesh = ref_process_mesh var_dist_attr . dims_mapping = [ - 1 ] self . _dist_context . set_tensor_dist_attr_for_program ( learning_var , var_dist_attr ) for input_name in op . desc . input_names ( ) : if input_name in [ 'Param' , 'Grad' , 'LearningRate' , "SkipUpdate" , "Beta1Tensor" , "Beta2Tensor" , "EpsilonTensor" , "MasterParam" ] : continue assert len ( op . desc . input ( input_name ) ) == 1 input_var = vars [ op . desc . input ( input_name ) [ 0 ] ] input_var_attr = TensorDistributedAttribute ( ) if "Beta1Pow" in input_name or "Beta2Pow" in input_name : input_var_attr . dims_mapping = [ - 1 ] op_dist_attr . set_input_dims_mapping ( input_var . name , [ - 1 ] ) op_dist_attr . set_output_dims_mapping ( input_var . name , [ - 1 ] ) else : assert "Moment" in input_name input_var_attr . dims_mapping = ref_dims_mapping op_dist_attr . set_input_dims_mapping ( input_var . name , ref_dims_mapping ) op_dist_attr . set_output_dims_mapping ( input_var . name , ref_dims_mapping ) input_var_attr . process_mesh = ref_process_mesh self . _dist_context . set_tensor_dist_attr_for_program ( input_var , input_var_attr ) self . _dist_context . set_op_dist_attr_for_program ( op , op_dist_attr ) continue