Introduction

If you want cosine similarity between every row of matrix A and every row of matrix B, the clean solution is to normalize both matrices row-wise and then use matrix multiplication. That gives you the full pairwise similarity matrix efficiently.

This is better than looping over rows in Python because PyTorch can do the whole operation in optimized tensor code.

The Shape You Want

Suppose:

• 'A has shape (m, d)'
• 'B has shape (n, d)'

You want an output of shape (m, n) where entry (i, j) is the cosine similarity between row i of A and row j of B.

Normalize Then Multiply

1import torch
2import torch.nn.functional as F
3
4A = torch.tensor([[1.0, 0.0], [1.0, 1.0]])
5B = torch.tensor([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]])
6
7A_norm = F.normalize(A, p=2, dim=1)
8B_norm = F.normalize(B, p=2, dim=1)
9
10similarity = A_norm @ B_norm.T
11print(similarity)

That final matrix multiplication is the key step.

Cosine similarity is just the dot product of normalized vectors, so once each row has unit length, A_norm @ B_norm.T produces all pairwise cosine values.

Why F.cosine_similarity Alone Is Not Enough

torch.nn.functional.cosine_similarity is great when you already have tensors aligned for one-to-one comparison along a given dimension.

For example:

x = torch.tensor([[1.0, 0.0], [1.0, 1.0]])
y = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
print(F.cosine_similarity(x, y, dim=1))

That compares row 0 with row 0 and row 1 with row 1.

It does not directly produce the full (m, n) all-pairs matrix unless you add broadcasting tricks. For full pairwise results, normalization plus matrix multiplication is usually the clearest answer.

Numerical Stability

F.normalize already handles normalization carefully, but the basic concern remains: zero vectors do not have a meaningful cosine similarity because their norm is zero.

If your data may contain zero rows, decide what similarity value you want in those cases. Often the safest approach is to avoid zero embeddings or filter them before similarity computation.

Batched and Large-Matrix Considerations

For moderate sizes, the pairwise matrix method is ideal. For very large m and n, the full (m, n) output may be too large to materialize in memory.

In that case, compute in blocks.

1def pairwise_cosine_blocks(A, B, block_size=1024):
2    A_norm = F.normalize(A, p=2, dim=1)
3    B_norm = F.normalize(B, p=2, dim=1)
4    blocks = []
5    for start in range(0, A_norm.size(0), block_size):
6        stop = start + block_size
7        blocks.append(A_norm[start:stop] @ B_norm.T)
8    return torch.cat(blocks, dim=0)

The math is identical. You are only trading memory for multiple smaller multiplications.

Common Pitfalls

The biggest mistake is using F.cosine_similarity expecting it to generate all pairwise comparisons automatically. By default it compares aligned slices, not all combinations.

Another common issue is forgetting to normalize the rows before matrix multiplication. Raw dot products are not cosine similarity.

People also run into memory problems when m and n are large and the full output matrix is huge.

Finally, zero vectors can make cosine-based reasoning unstable or undefined, so check your inputs if the output looks strange.

Summary

• Normalize rows of both matrices first.
• Compute pairwise cosine similarity with A_norm @ B_norm.T.
• This yields an (m, n) matrix for all row pairs.
• 'F.cosine_similarity is better for one-to-one aligned comparisons.'
• Use block processing when the full output matrix is too large.
• Be careful with zero vectors because cosine similarity is not well defined for them.