Given two 2D matrices of shape N x D, compute cosine similarity between A[i, :] and B[i, :]

rowwise_cosine_similarity.py

import numpy as np
import scipy.sparse

def cosine_similarity_kernel(
        matrix_a: Union[np.array, scipy.sparse.csr_matrix],
        matrix_b: Union[np.array, scipy.sparse.csr_matrix],
        eps: float = 1e-9,
) -> np.array:
    if scipy.sparse.issparse(matrix_a) or scipy.sparse.issparse(matrix_b):
        a_norm = (scipy.linalg.norm(matrix_a, axis=1) + eps).flatten()
        b_norm = (scipy.linalg.norm(matrix_b, axis=1) + eps).flatten()
        return np.asarray(matrix_a.multiply(matrix_b).sum(axis=1).flatten() / a_norm / b_norm)
    else:
        a_norm = np.linalg.norm(matrix_a, axis=1) + eps
        b_norm = np.linalg.norm(matrix_b, axis=1) + eps
        return np.sum(matrix_a * matrix_b, axis=1) / a_norm / b_norm