How to express tensor contraction efficiently? #einsum

autograd
1

I have 2 tensors of the following dimensions:

A: n x i x o
B: n x b x i

and I would like to compute the tensor C of dimension n x b x o. Here, n denotes the number of feature maps, o is the output dimension, i is the input dimension, and b is the batch size.

Think of A, B, C as stacks of matrices. The operation I’m looking for is essentially map-wise matrix multiplies.

What would be the most GPU-efficient way to express my computations.

Would

C = torch.einsum('nio,nbi->nbo', [A, B])

do the trick?

Would that be correct and reasonably efficient? If not, what’s a better alternative?

Note that I can change the orders of the dimensions to make the computation more efficient if necessary.

2

This particular form is a batch matrix multiplication if you swap A and B, so you could use torch.bmm or torch.matmul directly.
Einsum will (currently) reduce to bmm, so it’ll be similar in terms of performance, but with a few permutations.

Best regards

Thomas