Alternative to einsum for 3D tensor product?

inglada: <p>Hi,</p>
<p>I am doing temporal resampling of multivariate time series using a tensor product.</p>
<p>The goal is to transform an input signal which has C channels and Ti time stamps to an output signal with the same number of channels and To time stamps. Each channel is interpolated independently and the output time stamps are a simple linear combination of the input ones.</p>
<p>That is, the time stamp i of an output channel c is<br>
$y_c[i] = \sum_{l} a_c[i,l]x_c[l]$<br>
where x_c is the is the corresponding input channel.</p>
<p>This is of course just a matrix multiplication<br>
$y = x A$<br>
where A is the matrix containing the a_c[i,l] coefficients and y and x are vectors.</p>
<p>To do this with minibatches, a need a tensor product instead of a matrix<br>
product, since x will be a tensor with shape (batch,channels,in_dates)<br>
and A will have shape (batch,in_dates,out_dates) and I want the output<br>
to have shape (batch, channels, out_dates).</p>
<p>I have implemented this as follows (b is for Batch, c for Channels, i for Input time stamps and o for Output time stamps):<br>
y = torch.einsum(‘bci,bio-&gt;bco’, x, A)</p>
<p>The problem I am facing is that this is very slow. I guess that building the operation from a string does not allow any optimization and I was wondering if there is a way to implement this using other faster operations. Maybe there is some reshaping, (un)squeezing and broadcasting black magic, but I can’t figure it out!</p>
<p>Thank you for your help.</p>

KFrank: <p>Hi Jordi!</p>
<aside class="quote no-group" data-username="inglada" data-post="1" data-topic="115995" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/i/ba9def/48.png" class="avatar"> inglada:</div>
<blockquote>
<p>y = torch.einsum(‘bci,bio-&gt;bco’, x, A)</p>
<p>The problem I am facing is that this is very slow.</p>
</blockquote>
</aside>
<p>First let me note that in my (limited) experience <code>einsum()</code> is not<br>
particularly slow.  When I’ve timed it on some toy problems it’s<br>
been more-or-less as fast as alternatives based on reshaping the<br>
tensors being multiplied (and sometimes faster).  Having said that,<br>
I do recall seeing some github issues that suggest that <code>einsum()</code><br>
might better at optimizing some index-contractions than others.</p>
<aside class="quote no-group">
<blockquote>
<p>Maybe there is some reshaping, (un)squeezing and broadcasting black magic,</p>
</blockquote>
</aside>
<p>Yes, you can do this with a couple of <code>unsqueeze()</code>s to line things up<br>
right and then use <code>matmul()</code> with its “batch-dimensions” broadcasting<br>
to broadcast the absent channel dimension of <code>A</code>:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'1.7.1'
&gt;&gt;&gt; _ = torch.manual_seed (2021)
&gt;&gt;&gt; nBatch = 2
&gt;&gt;&gt; nChannel = 2
&gt;&gt;&gt; nInput = 3
&gt;&gt;&gt; nOutput = 5
&gt;&gt;&gt; x = torch.randn ((nBatch, nChannel, nInput))
&gt;&gt;&gt; A = torch.randn ((nBatch, nInput, nOutput))
&gt;&gt;&gt; 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]]])
&gt;&gt;&gt; torch.einsum ('bci, bio -&gt; 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]]])
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

inglada: <p>Hi,<br>
Thanks for your help!</p>
<aside class="quote no-group" data-username="KFrank" data-post="2" data-topic="115995">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/k/ecb155/48.png" class="avatar"> KFrank:</div>
<blockquote>
<p>I do recall seeing some github issues that suggest that <code>einsum()</code><br>
might better at optimizing some index-contractions than others.</p>
</blockquote>
</aside>
<p>Yes, I did some research before posting and saw these issues and some posts on this forum too. Maybe I could reorder the dimensions and time different options.</p>
<p>Your answer is the kind of approach I was looking for. Actually, I already have some reshaping earlier in the code because channels and time are originally interleaved in the data source, so I wonder if chaining all these operations might impact performace.</p>
<p>I will time it for my particular case and see if this speeds things up.</p>
<p>Thanks again!</p>