Calculate angles in pytorch

Question

If we have a set of points Rs, we can use torch.cdist to get the all pair distances.

dists_ij = torch.cdist(Rs, Rs)

Is there a function to get the angles between two set of vectors Vs like this:

angs_ij = torch.angs(Vs, Vs)

Answer 1

You can do this manually using the relation between the dot product of two vectors and the angle between them:

# normalize the vectors
nVs = Vs / torch.norm(Vs, p=2, dim=-1, keepdim=True)
# compute cosine of the angles using dot product
cos_ij = torch.einsum('bni,bmi->bnm', nVs, nVs)