docs: def forward(self, query, context):
          """
          Args:
              query (:class:`torch.FloatTensor` [batch size, output length, dimensions]): Sequence of
                  queries to query the context.
              context (:class:`torch.FloatTensor` [batch size, query length, dimensions]): Data
                  overwhich to apply the attention mechanism.
  
          Returns:
              :class:`tuple` with `output` and `weights`:
              * **output** (:class:`torch.LongTensor` [batch size, output length, dimensions]):
                Tensor containing the attended features.
              * **weights** (:class:`torch.FloatTensor` [batch size, output length, query length]):
                Tensor containing attention weights.
          """
          batch_size, output_len, dimensions = query.size()
          query_len = context.size(1)
  
          if self.attention_type == "general":
              query = query.view(batch_size * output_len, dimensions)
              query = self.linear_in(query)
              query = query.view(batch_size, output_len, dimensions)
  
          # TODO: Include mask on PADDING_INDEX?
  
          # (batch_size, output_len, dimensions) * (batch_size, query_len, dimensions) ->
          # (batch_size, output_len, query_len)
          attention_scores = torch.bmm(query, context.transpose(1, 2).contiguous())
  
          # Compute weights across every context sequence
          attention_scores = attention_scores.view(batch_size * output_len, query_len)
          attention_weights = self.softmax(attention_scores)
          attention_weights = attention_weights.view(batch_size, output_len, query_len)
  
          # (batch_size, output_len, query_len) * (batch_size, query_len, dimensions) ->
          # (batch_size, output_len, dimensions)
          mix = torch.bmm(attention_weights, context)
  
          # concat -> (batch_size * output_len, 2*dimensions)
          combined = torch.cat((mix, query), dim=2)
          combined = combined.view(batch_size * output_len, 2 * dimensions)
  
          # Apply linear_out on every 2nd dimension of concat
          # output -> (batch_size, output_len, dimensions)
          output = self.linear_out(combined).view(batch_size, output_len, dimensions)
          output = self.tanh(output)
  
          return output, attention_weights