TensorFlow | tf.math.top_k for getting indices and values

TensorFlow TensorFlow tf.math.top_k

This code snippet is written for TensorFlow2.0.

tf.math.top_k finds values and indices of specified number of large entries.

"tensors" in below code is a list of vectors, tf.keras.losses.cosine_similarity is used for calculating similarity between vectors, from the similarity vector tf.math.top_k provides values and indices of specified number of large entries.


import tensorflow as tf

tensors = [[.56, .89, .89, .58], [.14, .35, .67, .65],
            [.54, .95, .07, .5], [.64, .75, .81, .05],
            [.21, .67, .26, .87], [.63, .34, .89, .93]]

Similarity of one vector with all other vectors


co_similarity = tf.keras.losses.cosine_similarity(
    tensors[0],
      tensors,
    axis=-1
)
print(co_similarity)

====== Output =======
tf.Tensor([-1.  -0.90591335 -0.8337847  -0.93034923 -0.8450875  -0.902475  ], 
shape=(6,), dtype=float32)
tf.Tensor(-0.90591335, shape=(), dtype=float32)

To get two most similar vectors to vector 1 (tensor[0]), we need to get least values from co_similarity, these values can be fetched by tf.math.top_k , tf.math.top_k returns values and indices.



output = tf.math.top_k(co_similarity,k=len(co_similarity))

print(output.values[-2:])

====== Output ======
tf.Tensor([-0.93034923 -1.        ], shape=(2,), dtype=float32)


print(output.indices[-2:])
======= Output =======
tf.Tensor([3 0], shape=(2,), dtype=int32)