Alternative to einsum

PytorchLearning: <p>Hi all,</p>
<p>I try to do a capsule neural network (CapsNet) model from scratch to try to understand how it works</p>
<p>with this line of code it works</p>
<p>def input_caps2U(self, x):<br>
return torch.einsum(‘bij,ijkl-&gt;bikl’, x, self.weight).contiguous()</p>
<p>but i would like to replace einsum, does anyone have an idea how to do this and can they please show me because no matter what i change i always get an error.</p>
<p>Thank you for your help.</p>

KFrank: <p>Hi Py!</p>
<aside class="quote no-group" data-username="PytorchLearning" data-post="1" data-topic="178181" 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/p/d2c977/48.png" class="avatar"> PytorchLearning:</div>
<blockquote>
<p>i would like to replace einsum</p>
</blockquote>
</aside>
<p>With appropriate <code>permute()</code>s and <code>unsqueeze()</code>s to get the dimensions<br>
aligned properly, you can use (batch) <code>matmul()</code>.</p>
<p>Consider:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
2.0.0
&gt;&gt;&gt;
&gt;&gt;&gt; _ = torch.manual_seed (2023)
&gt;&gt;&gt;
&gt;&gt;&gt; b = 3
&gt;&gt;&gt; i = 4
&gt;&gt;&gt; j = 5
&gt;&gt;&gt; k = 2
&gt;&gt;&gt; l = 2
&gt;&gt;&gt;
&gt;&gt;&gt; x = torch.randn (b, i, j)
&gt;&gt;&gt; weight = torch.randn (i, j, k, l)
&gt;&gt;&gt;
&gt;&gt;&gt; resultA = torch.einsum ('bij, ijkl -&gt; bikl', x, weight)
&gt;&gt;&gt; xB = x.unsqueeze (2)
&gt;&gt;&gt; weightB = weight.permute (2, 3, 0, 1).unsqueeze (-1).unsqueeze (2)
&gt;&gt;&gt; resultB = (xB @ weightB).squeeze().permute (2, 3, 0, 1)
&gt;&gt;&gt;
&gt;&gt;&gt; torch.equal (resultA, resultB)
True
&gt;&gt;&gt;
&gt;&gt;&gt; resultA.shape
torch.Size([3, 4, 2, 2])
</code></pre>
<p>As an aside, <code>einsum()</code> works well.  Why would you want to forgo the<br>
convenience it offers?</p>
<p>Best.</p>
<p>K. Frank</p>

PytorchLearning: <p>Hi Frank,</p>
<p>First, thanks a lot for your help.</p>
<p>I would like to replace it to understand a little more how einsum works and to compare its performance and to see if there are other ways to optimize memory or make the calculation faster. By the way, I wanted to ask you three questions:</p>
<p>1- Could you please explain to me the code you have written to understand how you have managed to achieve the lines you have suggested, which works very well because I’ve been trying to find an alternative to einsum for a week.</p>
<p>2- Is there a way of comparing their performance (memory, speed, etc.)?</p>
<p>3- How I can use matmul() or Matrix Multiplication on Polars python only to learn how einsum works.</p>
<p>Thank you for your help.</p>

KFrank: <p>Hi Py!</p>
<aside class="quote no-group" data-username="PytorchLearning" data-post="4" data-topic="178181" 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/p/d2c977/48.png" class="avatar"> PytorchLearning:</div>
<blockquote>
<p>1- Could you please explain to me the code you have written</p>
</blockquote>
</aside>
<p>Print out the shapes of the relevant tensors step by step after each<br>
<code>permute()</code>, <code>unsqueeze()</code>, and <code>squeeze()</code> to see how the dimensions<br>
move around to get aligned.  Look at the documentation for <code>matmul()</code>.<br>
The key point is that if the two tensors both have more than two dimensions,<br>
then “batch” matrix multiplication is performed where matrix multiplication<br>
is performed on the last two dimensions, while the leading dimensions<br>
are treated as batch dimensions and must be broadcastable.  <code>weight</code>’s<br>
<code>unsqueeze (2)</code> is there to add a singleton dimension that gets broadcast<br>
against <code>x</code>’s <code>b</code> dimension.</p>
<aside class="quote no-group">
<blockquote>
<p>2- Is there a way of comparing their performance (memory, speed, etc.)?</p>
</blockquote>
</aside>
<p>Memory-profiling pytorch tends to be difficult, but I wouldn’t expect the<br>
two approaches to differ significantly in memory.  (However, broadcasting<br>
can sometimes use memory unexpectedly and seemingly unnecessarily.)</p>
<p>You can perform timing tests using python’s <code>time.time()</code>.  In earlier versions<br>
of pytorch we have sometimes seen <code>einsum()</code> perform unexpectedly and<br>
unnecessarily poorly in comparison to <code>matmul()</code> (and sometimes even in<br>
comparison to loops), and sometimes seen it outperform <code>matmul()</code><br>
(presumably due to some performance bug in <code>matmul()</code>).</p>
<p>Best.</p>
<p>K. Frank</p>

PytorchLearning: <p>Hi Frank,</p>
<p>Thanks a lot for the explanation, now I have a better understanding.</p>
<p>Best.</p>

J_Johnson: <p>The point of einsum is not necessarily for speed. In theory, they should be relatively the same, for most circumstances. (Granted, there may be cases where it works faster or slower.)</p>
<p>The reason some people prefer einsum, whether it be in Tensorflow, PyTorch or NumPy, is that it provides a more concise and natural way to write and understand the tensor operations.</p>
<p>You can think of einsum as a type of <strong>code shorthand for tensor operations</strong> that is universal between all machine learning libraries. Once you learn it, the tensor ops become easier to visualize, check, and troubleshoot, regardless of which library you are working in.</p>

PytorchLearning: <p>Hi J_Johnson,</p>
<p>Thanks a lot for these explanations which enabled me to learn a little more about einsum.</p>