google-research

1
# coding=utf-8
2
# Copyright 2024 The Google Research Authors.
3
#
4
# Licensed under the Apache License, Version 2.0 (the "License");
5
# you may not use this file except in compliance with the License.
6
# You may obtain a copy of the License at
7
#
8
#     http://www.apache.org/licenses/LICENSE-2.0
9
#
10
# Unless required by applicable law or agreed to in writing, software
11
# distributed under the License is distributed on an "AS IS" BASIS,
12
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
# See the License for the specific language governing permissions and
14
# limitations under the License.
15

16
"""Utilities for negative cache training."""
17

18
import tensorflow.compat.v2 as tf
19

20

21
def approximate_top_k_with_indices(negative_scores, k):
22
  """Approximately mines the top k highest scoreing negatives with indices.
23

24
  This function groups the negative scores into num_negatives / k groupings and
25
  returns the highest scoring element from each group. It also returns the index
26
  where the selected elements were found in the score matrix.
27

28
  Args:
29
    negative_scores: A matrix with the scores of the negative elements.
30
    k: The number of negatives to mine.
31

32
  Returns:
33
    The tuple (top_k_scores, top_k_indices), where top_k_indices describes the
34
    index of the mined elements in the given score matrix.
35
  """
36
  bs = tf.shape(negative_scores)[0]
37
  num_elem = tf.shape(negative_scores)[1]
38
  batch_indices = tf.range(num_elem)
39
  indices = tf.tile(tf.expand_dims(batch_indices, axis=0), multiples=[bs, 1])
40
  grouped_negative_scores = tf.reshape(negative_scores, [bs * k, -1])
41
  grouped_batch_indices = tf.range(tf.shape(grouped_negative_scores)[0])
42
  grouped_top_k_scores, grouped_top_k_indices = tf.math.top_k(
43
      grouped_negative_scores)
44
  grouped_top_k_indices = tf.squeeze(grouped_top_k_indices, axis=1)
45
  gather_indices = tf.stack([grouped_batch_indices, grouped_top_k_indices],
46
                            axis=1)
47
  grouped_indices = tf.reshape(indices, [bs * k, -1])
48
  grouped_top_k_indices = tf.gather_nd(grouped_indices, gather_indices)
49
  top_k_indices = tf.reshape(grouped_top_k_indices, [bs, k])
50
  top_k_scores = tf.reshape(grouped_top_k_scores, [bs, k])
51
  return top_k_scores, top_k_indices
52

53

54
def cross_replica_concat(tensor, axis):
55
  replica_context = tf.distribute.get_replica_context()
56
  return replica_context.all_gather(tensor, axis=axis)
57