CODEWITHSUNDEEP
pythonpytorchdistance
apgsov
apgsov

Reputation: 904

Understanding broadcasting and arithmetic operations on different dimension tensors

I'm currently working on computing various similarity metrics between vectors such as cosine similarity, euclidean distance, mahalanobis distance, etc. As I'm working with vectors that can be very large, I need compute time to be minimal.

I'm struggling to understand how to work with vectors of different dimensions (they, do, however, share one dimension) and how to work with this in PyTorch.

I have two vectors, A and B with dimensions [867, 768] and [621, 768], respectively.

I am trying to compute the following:

I'm aware that this is achieveable under the hood with the likes of scipy and numpy but I'm trying to avoid detaching and moving the vectors to the CPU for speed.

Can someone help me understand the logic of the operators required in PyTorch to achieve this?

Upvotes: 1

Views: 186

Answers (1)

Ivan
Ivan

Reputation: 40648

You could use fancy indexing on both input tensors to unsqueeze a dimension such that A and B have a shape of (1, 867, 768) and (621, 1, 768) respectively. The subtraction operation will then automatically broadcast the two tensors to identical shapes.

>>> diff = A[None]-B[:,None]
>>> diff.shape
torch.Size([621, 867, 768])

This is the typical approach when implementing batched-pairwise distances.

More specifically, notice the difference between the two variants: A[None]-B[:,None] and A[:,None]-B[None].

diff = A[None]-B[:,None]    # (1, k_a, b) - (k_b, 1, b) -> (k_b, k_a, b) - (k_b, k_a, b)
diff.shape                  # (k_b, k_a, b)

Compared to:

diff = A[:,None]-B[None]    # (k_a, 1, b) - (1, k_b, b) -> (k_a, k_b, b) - (k_a, k_b, b)
diff.shape                  # (k_a, k_b, b)

You can read more about broadcasting on the NumPy documentation page.

Upvotes: 1

Related Questions