Implementation of cos similarity matrix [Pytorch, Tensorflow]

Cos (cosine) similarity is one of the methods of contrast learning (Contrastive Learning) such as SimCLR that is used as an index of similarity in the feature space.

Since it was implemented by Tensorflow and Pytorch respectively, record it as a memo. (For your reference)

Pytorch

# input_size is(batchsize*Number of dimensions)

def cosine_matrix(a, b):
    dot = torch.matmul(a, torch.t(b))
    norm = torch.matmul(torch.norm(a, dim=1).unsqueeze(-1), torch.norm(b, dim=0).unsqueeze(0))
    return dot / norm

Tensorflow

def cosine_matrix(a, b):
    a_normed, _ = tf.linalg.normalize(a, axis=-1)
    b_normed, _ = tf.linalg.normalize(b, axis=-1)
    matrix = tf.matmul(a_normed, b_normed, transpose_b=True)
    return matrix