How to add an activation function to the intermediate result of torch.matmul

LIUJUN: <p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/5/7/570aecbbddd9f28efed33d7fe700b07264040694.png" data-download-href="https://discuss.pytorch.org/uploads/default/570aecbbddd9f28efed33d7fe700b07264040694" title="matmul"><img src="https://discuss.pytorch.org/uploads/default/original/3X/5/7/570aecbbddd9f28efed33d7fe700b07264040694.png" alt="matmul" data-base62-sha1="cq0WEJmsIaLAC650L09gxIi0Y3a" width="690" height="233" data-dominant-color="F8F1F1"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">matmul</span><span class="informations">799×270 24.2 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

KFrank: <p>Hi Jun!</p>
<p>If I understand your question correctly, you wish to apply some function, <code>f()</code>,<br>
to each scalar product of the form <code>m * n</code> in your matrix multiplication before<br>
you add them together.</p>
<p>That is, for <code>A</code> with shape <code>[3, 2]</code> and <code>B</code> of shape <code>[2, 3]</code> your product matrix<br>
<code>A @ B</code> will have shape <code>[3, 3]</code>, where each element of <code>A @ B</code> will be a sum of<br>
two terms and you want to apply <code>f()</code> to each of those eighteen terms individually<br>
before adding them together.</p>
<p>I would recommend using <code>einsum()</code> to perform the scalar multiplications while<br>
“deferring” the additions, apply <code>f()</code> element-wise to the tensor resulting from<br>
<code>einsum()</code>, and then sum the results.</p>
<p>Consider:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'2.3.0'
&gt;&gt;&gt; _ = torch.manual_seed (2024)
&gt;&gt;&gt; a = torch.randn (3, 2)
&gt;&gt;&gt; b = torch.randn (2, 3)
&gt;&gt;&gt; a
tensor([[-0.0404,  1.7260],
        [-0.8140,  1.3722],
        [ 0.5060, -0.4823]])
&gt;&gt;&gt; b
tensor([[-0.7853,  0.6681, -0.4439],
        [ 0.1888,  0.5986,  0.6458]])
&gt;&gt;&gt; a @ b
tensor([[ 0.3577,  1.0062,  1.1326],
        [ 0.8983,  0.2775,  1.2476],
        [-0.4884,  0.0493, -0.5361]])
&gt;&gt;&gt; torch.einsum ('ij, jk -&gt; ik', a, b)
tensor([[ 0.3577,  1.0062,  1.1326],
        [ 0.8983,  0.2775,  1.2476],
        [-0.4884,  0.0493, -0.5361]])
&gt;&gt;&gt; torch.einsum ('ij, jk -&gt; ikj', a, b)
tensor([[[ 0.0317,  0.3259],
         [-0.0270,  1.0332],
         [ 0.0179,  1.1147]],

        [[ 0.6392,  0.2591],
         [-0.5438,  0.8214],
         [ 0.3614,  0.8862]],

        [[-0.3973, -0.0911],
         [ 0.3380, -0.2887],
         [-0.2246, -0.3115]]])
&gt;&gt;&gt; torch.einsum ('ij, jk -&gt; ikj', a, b).sum (-1)
tensor([[ 0.3577,  1.0062,  1.1326],
        [ 0.8983,  0.2775,  1.2476],
        [-0.4884,  0.0493, -0.5361]])
&gt;&gt;&gt; torch.einsum ('ij, jk -&gt; ikj', a, b).sigmoid().sum (-1)
tensor([[1.0887, 1.2308, 1.2575],
        [1.2190, 1.0618, 1.2975],
        [0.8792, 1.0120, 0.8668]])
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

LIUJUN: <p>Dear Frank:<br>
Thank you very much for your detailed help, I have solved the problem in my project.</p>
<p>Best wishes from China.<br>
Liu Jun</p>