def get_nearest_neighbour(source, reference):
  """Get the nearest neighbour for every vector in source from reference."""
  normed_reference = tf.nn.l2_normalize(reference, axis=-1)
  normed_source = tf.nn.l2_normalize(source, axis=-1)

  cosine_sim = tf.matmul(normed_source, normed_reference, transpose_b=True)

  # Calculate the nearest neighbours and their cosine similarity
  nearest_neighbour = tf.argmax(cosine_sim, axis=-1)
  nearest_neighbour_sim = tf.reduce_max(cosine_sim, axis=-1)

  return nearest_neighbour, nearest_neighbour_sim
Get the nearest neighbour for every vector in source from reference.
language/bert_extraction/steal_bert_classifier/embedding_perturbations/embedding_util.py