title: PyTorch - applying attention efficiently
body: <p>I have build a RNN language model with attention and I am creating context vector for every element of the input by attending all the previous hidden states <em>(only one direction)</em>.</p>
  
  <p>The most straight forward solution in my opinion is using a <em>for-loop</em> over the RNN output, such that each context vector is computed one after another. </p>
  
  <pre><code>import torch
  import torch.nn as nn
  import torch.nn.functional as F
  
  class RNN_LM(nn.Module):
      def __init__(self, hidden_size, vocab_size, embedding_dim=None, droprate=0.5):
          super().__init__()
          if not embedding_dim:
              embedding_dim = hidden_size
          self.embedding_matrix = nn.Embedding(vocab_size, embedding_dim)
  
          self.lstm = nn.LSTM(input_size=embedding_dim, hidden_size=hidden_size, batch_first=False)
          self.attn = nn.Linear(hidden_size, hidden_size)
          self.vocab_dist = nn.Linear(hidden_size, vocab_size)
          self.dropout = nn.Dropout(droprate)
  
      def forward(self, x):
          x = self.dropout(self.embedding_matrix(x.view(-1, 1)))
          x, states = self.lstm(x)
          #print(x.size())
          x = x.squeeze()
          content_vectors = [x[0].view(1, -1)]
          # for-loop over hidden states and attention
          for i in range(1, x.size(0)):
              prev_states = x[:i]
              current_state = x[i].view(1, -1)
  
              attn_prod = torch.mm(self.attn(current_state), prev_states.t())
              attn_weights = F.softmax(attn_prod, dim=1)
              context = torch.mm(attn_weights, prev_states)
              content_vectors.append(context)
  
          return self.vocab_dist(self.dropout(torch.cat(content_vectors)))
  </code></pre>
  
  <p><em>Note: The <code>forward</code> method here is only used for training.</em></p>
  
  <p>However this solution is not very efficient as the code is not well parallelizable with computing each context vector sequently. But since the context vectors are not dependent on each other, I wonder if there is a non-sequential way of calculating them.</p>
  
  <p><strong>So is there is a way to compute the context vectors without <em>for-loop</em> so that more of computation can be parallelized?</strong></p>