similarity.py

import numpy as np
import tensorflow as tf

def np_cosine_similarity(u, v):
  u = np.expand_dims(u, 1)
  n = np.sum(u * v, axis=2)
  d = np.linalg.norm(u, axis=2) * np.linalg.norm(v, axis=1)

  return n / d

@tf.function
def tf_cosine_similarity(u, v):
  u = tf.expand_dims(u, 1)

  n = tf.reduce_sum(u * v, axis=2)
  d = tf.linalg.norm(u, axis=2) * tf.linalg.norm(v, axis=1)

  return n / d

# Generate random data
x = np.random.rand(5, 5)
y = np.random.rand(1, 5)

print("x:", "\n", x)
print("y:", "\n", y)

# Print cosine similarity in NumPy and TensorFlow
print("np:", "\n", np_cosine_similarity(x, y))
print("tf:", "\n", tf_cosine_similarity(x, y).numpy())