Tensorflow: Equivalent of scikit-learn's pairwise_kernels in tensorflow

Question

I am trying to find cosine similarity between two set of embeddings. I calculated it using pairwise_kernels from scikit-learn but now I am trying to leverage GPU, hence I am looking for the Tensorflow equivalent of pairwise_kernels. This is what I did using scikit-learn:

sim_elements = pairwise_kernels(input_array, orig_array, metric = 'cosine', n_jobs = -1)

I searched online and came across tensorflow pairwise_distances but that's not what I am trying to achieve. I am trying to achieve one-to-many similairity than one-to-one similarity. Any help is appreciated!

Answer 1

If I understand your question correct, I think you could just use the CosineSimilarity Loss-Function from tf.keras.losses: https://www.tensorflow.org/api_docs/python/tf/keras/losses/CosineSimilarity

But take care, they swapped the output that it works as a Loss-Funtion. So -1 is maximum similarity and 1 minimal similarity.