Hi Jordi!



 inglada:

y = torch.einsum(‘bci,bio->bco’, x, A)
The problem I am facing is that this is very slow.


First let me note that in my (limited) experience einsum() is not
particularly slow.  When I’ve timed it on some toy problems it’s
been more-or-less as fast as alternatives based on reshaping the
tensors being multiplied (and sometimes faster).  Having said that,
I do recall seeing some github issues that suggest that einsum()
might better at optimizing some index-contractions than others.


Maybe there is some reshaping, (un)squeezing and broadcasting black magic,


Yes, you can do this with a couple of unsqueeze()s to line things up
right and then use matmul() with its “batch-dimensions” broadcasting
to broadcast the absent channel dimension of A:
>>> import torch
>>> torch.__version__
'1.7.1'
>>> _ = torch.manual_seed (2021)
>>> nBatch = 2
>>> nChannel = 2
>>> nInput = 3
>>> nOutput = 5
>>> x = torch.randn ((nBatch, nChannel, nInput))
>>> A = torch.randn ((nBatch, nInput, nOutput))
>>> torch.matmul (x.unsqueeze (dim = 2), A.unsqueeze (dim = 1)).squeeze()
tensor([[[ 3.2144,  0.5790,  0.3564, -0.5355,  1.5428],
         [-2.4341, -1.2235, -0.5440,  0.9202, -0.3586]],

        [[-0.3876, -3.7793,  1.8061,  3.0808,  1.4999],
         [-0.6005, -0.9385,  5.1433,  0.8513,  0.3969]]])
>>> torch.einsum ('bci, bio -> bco', x, A)
tensor([[[ 3.2144,  0.5790,  0.3564, -0.5355,  1.5428],
         [-2.4341, -1.2235, -0.5440,  0.9202, -0.3586]],

        [[-0.3876, -3.7793,  1.8061,  3.0808,  1.4999],
         [-0.6005, -0.9385,  5.1433,  0.8513,  0.3969]]])

Best.
K. Frank