Einsum doubt in formulation

Jay_Timbadia: <p>Hi,</p>
<p>I just wanna know, is there any difference in the output of einsum of below mentioned two formulation.</p>
<p>torch.einsum(“bhld,lrd-&gt;bhlr”, query_layer, positional_embedding)<br>
torch.einsum(“bhrd,lrd-&gt;bhlr”, query_layer, positional_embedding)</p>
<p>Any help is much appreciated!</p>
<p>Jay</p>

tom: <p>Yes, there is, as the third axis of the first input tensor is aligned with dfferent axes in the second input and output.<br>
If I want to find out about these things, I usually choose random inputs of all different dimensions, e.g.</p>
<pre><code class="lang-python">query_layer = torch.randn(2, 3, 4, 5) # b h l d
positional_embedding = torch.randn(4, 6, 5)  # l r d
</code></pre>
<p>with this, you see that only the first einsum will work.<br>
Now in terms of what it should be, it’s not entirely clear to me.<br>
Typical dimensions at work in attention layers are batch, head, query position, embedding position = key, and feature. If we guess <em>b</em>atch and <em>h</em>ead and feature <em>d</em>imension, we have <em>l</em> and <em>r</em> for query and key, but then it is very atypical to have lrd as the second factor and (if we assume l=query, r=key) it would be more common to have <code>hrd</code> or more commonly <code>bhrd</code> for the second factor.<br>
In my small toroidal library, I implemented a <a href="https://github.com/MathInf/toroidal/blob/bff09f725627e4629d464008dc7c5f9d6322ebad/toroidal/models.py#L187" rel="noopener nofollow ugc">standard (BERT, GPT, ViT, …) attention module with einsum</a>. I number the key and query positions s and t and the features c (for channel), and then have use <code>bthc,bshc-&gt;bhts</code> after ordering the dimensions.</p>
<p>Best regards</p>
<p>Thomas</p>

Jay_Timbadia: <p>Thank you!</p>
<p>Also one more einsum for fourier transform is kind of not working.</p>
<pre><code class="lang-auto">class FourierMMLayer(nn.Module):
    def __init__(self, config):
        super().__init__()

        self.dft_mat_seq = torch.tensor(linalg.dft(512))
        self.dft_mat_hidden = torch.tensor(linalg.dft(768))

    def forward(self, hidden_states):
        hidden_states_complex = hidden_states.type(torch.complex128)
         #pre fourier torch.Size([2, 9, 768]) of hidden states
        return torch.einsum(
            "...ij,...jk,...ni-&gt;...nk",
            hidden_states_complex,
            self.dft_mat_hidden,
            self.dft_mat_seq
        ).real.type(torch.float32)
</code></pre>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "inference.py", line 22, in &lt;module&gt;
    obj1.forward(input_ids, token_type_ids)
  File "/mnt/sda1/ml_models/fourier_net/fnet.py", line 124, in forward
    self.encoder(input_ids, type_ids)
  File "/mnt/sda1/luck/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/mnt/sda1/ml_models/fourier_net/fnet.py", line 113, in forward
    sequence_output = self.encoder(embedding_output)
  File "/mnt/sda1/luck/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/mnt/sda1/ml_models/fourier_net/fnet.py", line 94, in forward
    hidden_states = layer_module(hidden_states)
  File "/mnt/sda1/luck/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/mnt/sda1/ml_models/fourier_net/fnet.py", line 80, in forward
    fft_output = self.fft(hidden_states)
  File "/mnt/sda1/luck/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/mnt/sda1/ml_models/fourier_net/fnet.py", line 62, in forward
    return torch.einsum(
  File "/mnt/sda1/luck/lib/python3.8/site-packages/torch/functional.py", line 299, in einsum
    return _VF.einsum(equation, operands)  # type: ignore[attr-defined]
RuntimeError: einsum(): operands do not broadcast with remapped shapes [original-&gt;remapped]: [2, 9, 768]-&gt;[2, 1, 1, 9, 768] [768, 768]-&gt;[1, 1, 768, 1, 768] [512, 512]-&gt;[1, 512, 1, 512, 1]
</code></pre>
<p>Can you please help?</p>

tom: <p>You dimension i is 9 in the first input, and 512 in the last, that doesn’t seem right.</p>

Jay_Timbadia: <p>Shape is correct <a class="mention" href="/u/tom">@tom</a><br>
Also that is the input embedding shape [batch_size, seq_len, embeddings]<br>
512 is number of positions with 768 embeddings.<br>
That seems ok to me.</p>
<p>I have implemented from this source</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/erksch/fnet-pytorch/blob/master/fnet.py">
  <header class="source">

      <a href="https://github.com/erksch/fnet-pytorch/blob/master/fnet.py" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/erksch/fnet-pytorch/blob/master/fnet.py" target="_blank" rel="noopener nofollow ugc">erksch/fnet-pytorch/blob/master/fnet.py</a></h4>


      <pre><code class="lang-py">import re
import torch
import torch.utils.checkpoint
from scipy import linalg
from torch import nn


class FNetEmbeddings(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.word_embeddings = nn.Embedding(config['vocab_size'], config['embedding_size'],
                                            padding_idx=config['pad_token_id'])
        self.position_embeddings = nn.Embedding(config['max_position_embeddings'], config['embedding_size'])
        self.token_type_embeddings = nn.Embedding(config['type_vocab_size'], config['embedding_size'])

        self.layer_norm = nn.LayerNorm(config['embedding_size'], eps=config['layer_norm_eps'])
        self.hidden_mapping = nn.Linear(config['embedding_size'], config['hidden_size'])
        self.dropout = nn.Dropout(config['dropout_rate'])

        self.register_buffer(
</code></pre>



  This file has been truncated. <a href="https://github.com/erksch/fnet-pytorch/blob/master/fnet.py" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>


tom: <p>I am hesitant to insist, but so the error message literally means you have a shape mismatch.<br>
The best way to look into this would be that you print out the dimensions of your inputs and write a letter for each in the einsum formula next to it. Then it should get more clear what is going on.</p>
<p>Best regards</p>
<p>Thomas</p>