def greedy_updates(old_elements, grad_difference, biggest_indices,
                   max_seq_length):
  """Construct new beams with greedy updates based on gradient differences."""
  input_mask = np.array([elem["input_mask"] for elem in old_elements])
  # Mask tokens which cannot be flipped due to the input template constraint
  masked_gd = input_mask * np.reshape(grad_difference, [-1, max_seq_length])

  # Find the top flip score estimates by checking over tokens and word pieces
  top_indices = np.argmax(masked_gd, axis=1)
  top_scores = np.max(masked_gd, axis=1)

  new_elements = []
  for elem_num, (score, seq_num) in enumerate(zip(top_scores, top_indices)):
    old_elem = old_elements[elem_num]
    if old_elem["stopped"]:
      new_elements.append(old_elem)
    else:
      new_input_ids = list(old_elem["input_ids"])
      new_input_ids[seq_num] = biggest_indices[elem_num, seq_num, 0]
      new_elem = {
          "input_ids": np.array(new_input_ids),
          "original_input_ids": old_elem["original_input_ids"],
          "ip_num": old_elem["ip_num"],
          "score": score,
          "bert_input_mask": old_elem["bert_input_mask"],
          "input_mask": old_elem["input_mask"],
          "token_type_ids": old_elem["token_type_ids"],
          "prob_vector": old_elem["prob_vector"],
          "stopped": False,
          "steps_taken": old_elem["steps_taken"]
      }
      new_elements.append(new_elem)

  return new_elements
Construct new beams with greedy updates based on gradient differences.
language/bert_extraction/steal_bert_classifier/embedding_perturbations/discrete_invert_embeddings.py