Tricky einsum operation going wrong

matheussfarias: <p>Hi folks,</p>
<p>Consider the matrices A and B such that the shape of A is <code>torch.Size([4096, 25088, 8])</code> and the shape of B is <code>torch.Size([256, 25088])</code>. The final result I want to obtain is the matrix C resulted from the following <code>einsum</code> operation</p>
<pre><code class="lang-auto">C = torch.einsum('bij,ki-&gt;bkj', A, B)
</code></pre>
<p>However, to get C, I want to compute sections of A and B, operate them using <code>einsum</code> and then accumulating over the number of sections using <code>torch.sum</code>. Now, consider A_sec with shape <code>torch.Size([4096, 196, 128, 8])</code> and B_sec with shape <code>torch.Size([256, 196, 128])</code>. In this case, I have 196 sections of 128 vectors. The same result is obtained doing</p>
<pre><code class="lang-auto">C = torch.einsum('bijk,lij-&gt;blk', A_sec, B_sec)
</code></pre>
<p>But, as mentioned before, I want to accumulate on the number of sections, something like</p>
<pre><code class="lang-auto"># THIS PIECE OF CODE IS WRONG!
C_sec = torch.einsum('bijk,lij-&gt;bilk', A_sec, B_sec)
C = torch.sum(C_sec, axis=1)
</code></pre>
<p>But this answer doesn’t provide the C matrix as before. What is wrong with what I did? What operations do I need to perform using <code>einsum</code> to get the result I want?</p>
<p>Thanks</p>

KFrank: <p>Hi Matheus!</p>
<aside class="quote no-group" data-username="matheussfarias" data-post="1" data-topic="188843" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/matheussfarias/48/63894_2.png" class="avatar"> matheussfarias:</div>
<blockquote>
<p>But this answer doesn’t provide the C matrix as before.</p>
</blockquote>
</aside>
<p>Your two versions of <code>einsum()</code> work for me:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
2.0.1
&gt;&gt;&gt;
&gt;&gt;&gt; _ = torch.manual_seed (2023)
&gt;&gt;&gt;
&gt;&gt;&gt; nb = 4096
&gt;&gt;&gt; ni = 196
&gt;&gt;&gt; nj = 128
&gt;&gt;&gt; nk = 8
&gt;&gt;&gt; nl = 256
&gt;&gt;&gt;
&gt;&gt;&gt; A_sec = torch.randn (nb, ni, nj, nk)
&gt;&gt;&gt; B_sec = torch.randn (nl, ni, nj)
&gt;&gt;&gt;
&gt;&gt;&gt; A_sec.shape
torch.Size([4096, 196, 128, 8])
&gt;&gt;&gt; B_sec.shape
torch.Size([256, 196, 128])
&gt;&gt;&gt;
&gt;&gt;&gt; C1 = torch.einsum ('bijk, lij -&gt; blk', A_sec, B_sec)
&gt;&gt;&gt; C_sec = torch.einsum ('bijk, lij -&gt; bilk', A_sec, B_sec)
&gt;&gt;&gt; C2 = torch.sum (C_sec, axis = 1)
&gt;&gt;&gt;
&gt;&gt;&gt; torch.allclose (C1, C2, atol = 1.e-3)
True
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

matheussfarias: <p>Hi Frank!</p>
<p>This is very interesting, the way I was comparing C1 and C2 in your code was by writing</p>
<pre><code class="lang-auto">(C1 == C2).all()
</code></pre>
<p>In fact, when I do this in your code, it prints False even though it works with tol 1.e-3. I tried to 1.e-4, and it also starts to be False. In my previous examples, C1 and C2 were equal in <code>(C1 == C2).all()</code>. Do you know any reason why this is happening?</p>

KFrank: <p>Hi Matheus!</p>
<aside class="quote no-group" data-username="matheussfarias" data-post="3" data-topic="188843" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/matheussfarias/48/63894_2.png" class="avatar"> matheussfarias:</div>
<blockquote>
<pre><code class="lang-auto">(C1 == C2).all()
</code></pre>
<p>In fact, when I do this in your code, it prints False</p>
</blockquote>
</aside>
<p>This tests for exact equality, which you can’t, in general, expect in the presence<br>
of floating-point round-off error.  (The two versions sum their sums in different<br>
orders, so the results can differ by round-off error.)</p>
<p>The tensors being multiplied are quite large, so individual elements of the result<br>
range up to about <code>1.e3</code> in magnitude.  An absolute tolerance (<code>atol</code>) of <code>1.e-3</code><br>
then corresponds to about six decimal digits of agreement between the two results,<br>
which is consistent with the six-to-seven digits you expect from single-precision<br>
(<code>float</code>) arithmetic.</p>
<p>Best.</p>
<p>K. Frank</p>

matheussfarias: <p>Thanks a lot for the explanation! Really appreciate that <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>