Mixed dimension difference operation with einsum

I’m struggling with the einsum notation.

In the simple case where I have tensor A size: (1, 5) and I want to find the difference over a series of vectors B size: (6, 5) I am able to simply do something like result = A - B . For my batch training the dimensions increase. For example a batch of 3: A size: (3, 5) B size: (3, 6, 5). I now need to find the difference matrix for each vector in A relative to each matrix in B.

I believe there is probably a way to do this with torch.einsum() however I cannot wrap my head around the process. I’m able to successfully multiply them as I would expect using: torch.einsum('ij,ikj->ikj', [A, B]) but I cannot figure out how to write the statement to effectively perform the difference calculation. I’ve found a few resources for einsum documentation but I’m still a bit confused. Is this actually possible with einsum()? Or is there another option?

Hi Pumplerod!

Use broadcasting with an unsqueeze() to get the dimensions lined
up properly:

>>> import torch
>>> torch.__version__
>>> _ = torch.manual_seed (2022)
>>> A = torch.randn (3, 5)
>>> B = torch.randn (3, 6, 5)
>>> resultBroadcast = A.unsqueeze (1) - B
>>> resultBroadcast.shape
torch.Size([3, 6, 5])

But if you really want to use einsum() you can just convert addition
into multiplication:

>>> resultEinsum = torch.einsum('ij,ikj->ikj', A.exp(), (-B).exp()).log()
>>> torch.allclose (resultBroadcast, resultEinsum)


K. Frank

That is so helpful! Thank you. Of course, using log for the matrix multiplication makes sense. I have a lot to learn about einsum. I’ve only just discovered it, but looks super powerful.