docstr:     Attention function
  Parameters
  ----------
  query: torch.Tensor, shape=[B, D]
  memory: torch.Tensor, shape=[B, N, D]
  memory_mask: torch.Tensor, shape=[B, N]
  coverage: torch.Tensor, shape=[B, N, D]
  
  Returns
  -------
  attn_results: torch.Tensor, shape=[B, D]
  attn_scores: torch.Tensor, shape=[B, N]
docstr_tok: ["Attention","function","Parameters","--","--","--","--","--","query",":","torch.Tensor",",","shape=","[","B",",","D","]","memory",":","torch.Tensor",",","shape=","[","B",",","N",",","D","]","memory_mask",":","torch.Tensor",",","shape=","[","B",",","N","]","coverage",":","torch.Tensor",",","shape=","[","B",",","N",",","D","]","Returns","--","--","--","-","attn_results",":","torch.Tensor",",","shape=","[","B",",","D","]","attn_scores",":","torch.Tensor",",","shape=","[","B",",","N","]"]
code: def forward(self, query, memory, memory_mask=None, coverage=None):
      """
          Attention function
      Parameters
      ----------
      query: torch.Tensor, shape=[B, D]
      memory: torch.Tensor, shape=[B, N, D]
      memory_mask: torch.Tensor, shape=[B, N]
      coverage: torch.Tensor, shape=[B, N, D]
  
      Returns
      -------
      attn_results: torch.Tensor, shape=[B, D]
      attn_scores: torch.Tensor, shape=[B, N]
      """
      assert len(query.shape) == 2
      assert len(memory.shape) == 3
      assert query.shape[0] == memory.shape[0]
  
      aligns = self._calculate_aligns(query, memory, coverage=coverage)
  
      if memory_mask is not None:
          aligns = aligns.masked_fill(memory_mask == 0, -self.inf)
      scores = torch.softmax(aligns, dim=-1)
      scores = self.dropout(scores)
      ret = torch.bmm(scores.unsqueeze(1), memory)
      return ret.squeeze(1), scores
  
imports: import torch
  import torch.nn
label_desc: Use image node_id 2 for calling the Attention obj's underlying member method code with example usage: obj.forward(query, memory, memory_mask, coverage) and returns: ret,  scores