class_docstr: This class implements location-aware attention module for seq2seq learning.
  
  Reference: Attention-Based Models for Speech Recognition, Chorowski et.al.
  https://arxiv.org/pdf/1506.07503.pdf
  
  Arguments
  ---------
  attn_dim : int
      Size of the attention feature.
  output_dim : int
      Size of the output context vector.
  conv_channels : int
      Number of channel for location feature.
  kernel_size : int
      Kernel size of convolutional layer for location feature.
  scaling : float
      The factor controls the sharpening degree (default: 1.0).
  
  Example
  -------
  >>> enc_tensor = torch.rand([4, 10, 20])
  >>> enc_len = torch.ones([4]) * 10
  >>> dec_tensor = torch.rand([4, 25])
  >>> net = LocationAwareAttention(
  ...     enc_dim=20,
  ...     dec_dim=25,
  ...     attn_dim=30,
  ...     output_dim=5,
  ...     conv_channels=10,
  ...     kernel_size=100)
  >>> out_tensor, out_weight = net(enc_tensor, enc_len, dec_tensor)
  >>> out_tensor.shape
  torch.Size([4, 5])
docstr: Returns the output of the attention module.
  
  Arguments
  ---------
  enc_states : torch.Tensor
      The tensor to be attended.
  enc_len : torch.Tensor
      The real length (without padding) of enc_states for each sentence.
  dec_states : torch.Tensor
      The query tensor.
docstr_tok: ["Returns","the","output","of","the","attention","module",".","Arguments","--","--","--","--","-","enc_states",":","torch.Tensor","The","tensor","to","be","attended",".","enc_len",":","torch.Tensor","The","real","length","(","without","padding",")","of","enc_states","for","each","sentence",".","dec_states",":","torch.Tensor","The","query","tensor","."]
code: def forward(self, enc_states, enc_len, dec_states):
      """Returns the output of the attention module.
  
      Arguments
      ---------
      enc_states : torch.Tensor
          The tensor to be attended.
      enc_len : torch.Tensor
          The real length (without padding) of enc_states for each sentence.
      dec_states : torch.Tensor
          The query tensor.
      """
      if self.precomputed_enc_h is None:
          self.precomputed_enc_h = self.mlp_enc(enc_states)
          self.mask = length_to_mask(
              enc_len,
              max_len=enc_states.size(1),
              device=enc_states.device,
          )
  
          # multiply mask by 1/Ln for each row
          self.prev_attn = self.mask * (1 / enc_len.float()).unsqueeze(
              1
          )
  
      # compute location-aware features
      # [B, 1, L] -> [B, C, L]
      attn_conv = self.conv_loc(self.prev_attn.unsqueeze(1))
      # [B, C, L] -> [B, L, C] -> [B, L, F]
      attn_conv = self.mlp_loc(attn_conv.transpose(1, 2))
  
      dec_h = self.mlp_dec(dec_states.unsqueeze(1))
      attn = self.mlp_attn(
          torch.tanh(self.precomputed_enc_h + dec_h + attn_conv)
      ).squeeze(-1)
  
      # mask the padded frames
      attn = attn.masked_fill(self.mask == 0, -np.inf)
      attn = self.softmax(attn * self.scaling)
  
      # set prev_attn to current attn for the next timestep
      self.prev_attn = attn.detach()
  
      # compute context vectors
      # [B, 1, L] X [B, L, F]
      context = torch.bmm(attn.unsqueeze(1), enc_states).squeeze(1)
      context = self.mlp_out(context)
  
      return context, attn
  
imports: import torch
  import logging
  import torch.nn
  import numpy
  from typing import Optional
  from speechbrain.dataio.dataio import length_to_mask
  import torch.nn.functional
  import math
label_desc: Use image node_id 3 for calling the LocationAwareAttention obj's underlying member method code with example usage: obj.forward(enc_states, enc_len, dec_states) and returns: context,  attn