def _get_mcd_idx_1(divergence, examples_1, compounds_1, compounds_2, atoms,
                   get_compounds_fn, get_atoms_fn, start_idx):
  """Return index of example to swap from examples_1 to examples_2."""
  for example_idx, example in _shifted_enumerate(examples_1, start_idx):
    # Ensure example does not contain any atom that appears only once in
    # examples_1. Otherwise, we would violate the atom constraint.
    if _contains_atom(example, atoms, get_atoms_fn):
      continue

    # Compute the new compound divergence if we move the example from examples_1
    # to examples_2, ignoring the effect of moving some other example in
    # examples_2 to examples_1 for now.
    compounds_example = get_compounds_fn(example)
    new_divergence = _compute_new_divergence_1(compounds_1, compounds_2,
                                               compounds_example, divergence)

    # Return the first example that we find that would increase compound
    # divergence.
    if new_divergence > divergence:
      return example_idx, example
  return None, None
Return index of example to swap from examples_1 to examples_2.
language/compgen/nqg/tasks/mcd_utils.py