Matrix multipication shapes

username: ClaartjeBarkhof
cooked: <p>Hi there,</p>
<p>I would like to do a matrix multipication which I am not sure of how to implement. It would be an implementation of a doing a different linear forward for every 2D element in the batch.</p>
<p>For every 2D element of shape [seq_len, hidden_in] in the batch I would like to multiply with a specific matrix of shape [hidden_out, hidden_in] to create a batch output of elements of shape [seq_len, hidden_out]. So, the term batch here is a bit misleading as usually you would like to perform the same operation on all elements in the batch, but I would like to do a different operation for each element in the batch, but the same operation along the sequence dimension.</p>
<p>The (batched) shapes are as follows:</p>
<p>Input shape: <code>[batch, seq_len, hidden_in]</code><br>
Matrix: <code>[batch, hidden_out, hidden_in]</code><br>
Desired output shape: <code>[batch, seq_len, hidden_out]</code></p>
<p>It would replace an implementation with linear layers in a loop, something along the lines of:</p>
<pre><code class="lang-auto"># weights_for_elements_in_batch is of shape [batch, hidden_out, hidden_in]
# input is of shape [batch, seq_len, hidden_in]
outs = []
for i in range(batch_size):
    out = torch.nn.functional.linear(input[i, :, :], weight=weights_for_elements_in_batch[i, :, :])
output = torch.stack(outs, dim=0) # stack all elements in batch
</code></pre>
<p>If anyone has a suggestion, please let me know. Help is appreciated!</p>
<p>Cheers,</p>
<p>Claartje Barkhof</p>

username: rad
cooked: <p>Hey!</p>
<p>PyTorch provides the Einsum function, which allows you to do different tensor operations based on indices. See <a href="https://stackoverflow.com/questions/26089893/understanding-numpys-einsum#:~:text=How%20does%20einsum%20work%3F&amp;text=So%20here%2C%20the%20indexing%20operation,summed%20to%20return%20the%20answer." rel="noopener nofollow ugc">here</a>, <a href="https://stackoverflow.com/questions/55894693/understanding-pytorch-einsum" rel="noopener nofollow ugc">here</a>, <a href="https://www.youtube.com/watch?v=pkVwUVEHmfI" rel="noopener nofollow ugc">here</a> for more explanations. I find einsum the easiest way to do such operations.</p>
<p>(Also, I would assume that matrix is actually of shape [batch, hidden_in, hidden_out], as you want to multiply 2D slices of shape [seq_len, hidden_in] with 2D slices of shape [hidden_in, hidden_out], and that’s the only way that you could achieve that multiplication).</p>
<pre><code class="lang-auto">batch = 8
seq_len = 25
hidden_in = 16
hidden_out = 32

input = torch.randn((batch, seq_len, hidden_in))
matrix = torch.randn((batch, hidden_in, hidden_out))

output = torch.einsum('bsi, bio -&gt; bso', input, matrix).
output.shape

torch.Size([8, 25, 32]).
</code></pre>
<p>The most important element is the formula ‘bsi, bio → bso’, which basically says:</p>
<p>I have two inputs with shape indices (b, s, i), (b, i o) - hopefully you see the meaning of each letter, and an output of shape (b, s, o). I multiply corresponding 2d slices on the dimension b (meaning slices of shape (s, i) with slices of shape (i, o)). These “si, io → so” is basically regular 2D matrix multiplication rule, where you multiply and sum along the dimension i, common to each operand, and get the result of shape (s, o), for each pair in the batch dimension.</p>

username: ClaartjeBarkhof
cooked: <p>Thanks a lot <a class="mention" href="/u/rad">@rad</a> . That’s very informative &amp; helpful.</p>