def beam_search(old_beams, grad_difference, biggest_indices, beam_size,
                accumulate_scores, max_seq_length):
  """Construct new beams using the old ones based on gradient differences."""
  input_mask = np.array([beam["input_mask"] for beam in old_beams])
  # Mask tokens which cannot be flipped due to the input template constraint
  masked_gd = input_mask * np.reshape(grad_difference, [-1, max_seq_length])
  if accumulate_scores:
    # Update scores to include old scores in previous beam
    # old_scores.shape = [beam_size, 1, 1], beam_size = 1 for first iteration
    old_scores = np.array([[[x["score"]]] for x in old_beams])
    updated_gd_scores = masked_gd + old_scores
  else:
    updated_gd_scores = masked_gd
  # Find the top flip score estimates by checking over tokens and word pieces
  masked_flat_gd = np.reshape(updated_gd_scores, -1)
  top_flat_indices = np.argsort(masked_flat_gd)[::-1][:beam_size]
  top_indices = np.unravel_index(top_flat_indices, dims=masked_gd.shape)
  top_scores = masked_flat_gd[top_flat_indices]

  new_beams = []
  for score, beam_num, seq_num, topk_num in zip(top_scores, top_indices[0],
                                                top_indices[1], top_indices[2]):
    old_beam = old_beams[beam_num]
    if old_beam["stopped"]:
      new_beams.append(old_beam)
    else:
      new_input_ids = list(old_beam["input_ids"])
      new_input_ids[seq_num] = biggest_indices[beam_num, seq_num, topk_num]
      new_beam = {
          "input_ids": np.array(new_input_ids),
          "original_input_ids": old_beam["original_input_ids"],
          "ip_num": old_beam["ip_num"],
          "score": score,
          "bert_input_mask": old_beam["bert_input_mask"],
          "input_mask": old_beam["input_mask"],
          "token_type_ids": old_beam["token_type_ids"],
          "prob_vector": old_beam["prob_vector"],
          "stopped": False,
          "steps_taken": old_beam["steps_taken"]
      }
      new_beams.append(new_beam)

  return new_beams
Construct new beams using the old ones based on gradient differences.
language/bert_extraction/steal_bert_classifier/embedding_perturbations/discrete_invert_embeddings.py