docstr: Attention mechanism for radtts. Unlike in Flowtron, we have no
  restrictions such as causality etc, since we only need this during
  training.
  
  Args:
      queries (torch.tensor): B x C x T1 tensor (likely mel data)
      keys (torch.tensor): B x C2 x T2 tensor (text data)
      query_lens: lengths for sorting the queries in descending order
      mask (torch.tensor): uint8 binary mask for variable length entries
                           (should be in the T2 domain)
  Output:
      attn (torch.tensor): B x 1 x T1 x T2 attention mask.
                           Final dim T2 should sum to 1
docstr_tok: ["Attention","mechanism","for","radtts",".","Unlike","in","Flowtron",",","we","have","no","restrictions","such","as","causality","etc",",","since","we","only","need","this","during","training",".","Args",":","queries","(","torch.tensor",")",":","B","x","C","x","T1","tensor","(","likely","mel","data",")","keys","(","torch.tensor",")",":","B","x","C2","x","T2","tensor","(","text","data",")","query_lens",":","lengths","for","sorting","the","queries","in","descending","order","mask","(","torch.tensor",")",":","uint8","binary","mask","for","variable","length","entries","(","should","be","in","the","T2","domain",")","Output",":","attn","(","torch.tensor",")",":","B","x","1","x","T1","x","T2","attention","mask",".","Final","dim","T2","should","sum","to","1"]
code: def forward(
      self,
      queries,
      keys,
      query_lens,
      mask=None,
      key_lens=None,
      attn_prior=None,
  ):
      """Attention mechanism for radtts. Unlike in Flowtron, we have no
      restrictions such as causality etc, since we only need this during
      training.
  
      Args:
          queries (torch.tensor): B x C x T1 tensor (likely mel data)
          keys (torch.tensor): B x C2 x T2 tensor (text data)
          query_lens: lengths for sorting the queries in descending order
          mask (torch.tensor): uint8 binary mask for variable length entries
                               (should be in the T2 domain)
      Output:
          attn (torch.tensor): B x 1 x T1 x T2 attention mask.
                               Final dim T2 should sum to 1
      """
      temp = 0.0005
      keys_enc = self.key_proj(keys)  # B x n_attn_dims x T2
      # Beware can only do this since query_dim = attn_dim = n_mel_channels
      queries_enc = self.query_proj(queries)
  
      # Gaussian Isotopic Attention
      # B x n_attn_dims x T1 x T2
      attn = (queries_enc[:, :, :, None] - keys_enc[:, :, None]) ** 2
  
      # compute log-likelihood from gaussian
      eps = 1e-8
      attn = -temp * attn.sum(1, keepdim=True)
      if attn_prior is not None:
          attn = self.log_softmax(attn) + torch.log(
              attn_prior[:, None] + eps
          )
  
      attn_logprob = attn.clone()
  
      if mask is not None:
          attn.data.masked_fill_(
              mask.permute(0, 2, 1).unsqueeze(2), -float("inf")
          )
  
      attn = self.softmax(attn)  # softmax along T2
      return attn, attn_logprob
  
imports: from typing import Optional, Tuple
  import numpy
  import torch
  from torch import Tensor, nn
  from torch.cuda import amp
  from torch.cuda.amp import autocast
  from torch.nn import functional
  from nemo.collections.tts.modules.submodules import ConvNorm, LinearNorm, MaskedInstanceNorm1d
  from nemo.collections.tts.parts.utils.helpers import get_mask_from_lengths, sort_tensor, unsort_tensor
  from nemo.collections.tts.parts.utils.splines import piecewise_linear_inverse_transform, piecewise_linear_transform, unbounded_piecewise_quadratic_transform
label_desc: Use image node_id 2 for calling the ConvAttention obj's underlying member method code with example usage: obj.forward(queries, keys, query_lens, mask, key_lens, attn_prior) and returns: attn,  attn_logprob