class_docstr: Self-attention for lm tokens and text. 
class_docstr_tok: ["Self-attention","for","lm","tokens","and","text","."]
docstr: Forward pass of self-attention.
  
  Args:
      queries (torch.tensor): B x T1 x C1 tensor
      keys (torch.tensor): B x T2 x C2 tensor
      values (torch.tensor): B x T2 x C2 tensor
      q_mask (torch.tensor): B x T1 tensor, bool mask for variable length entries
      kv_mask (torch.tensor): B x T2 tensor, bool mask for variable length entries
  Output:
      attn_out (torch.tensor): B x T1 x C1 tensor
docstr_tok: ["Forward","pass","of","self-attention",".","Args",":","queries","(","torch.tensor",")",":","B","x","T1","x","C1","tensor","keys","(","torch.tensor",")",":","B","x","T2","x","C2","tensor","values","(","torch.tensor",")",":","B","x","T2","x","C2","tensor","q_mask","(","torch.tensor",")",":","B","x","T1","tensor",",","bool","mask","for","variable","length","entries","kv_mask","(","torch.tensor",")",":","B","x","T2","tensor",",","bool","mask","for","variable","length","entries","Output",":","attn_out","(","torch.tensor",")",":","B","x","T1","x","C1","tensor"]
code: def forward(self, queries, keys, values, q_mask=None, kv_mask=None):
      """Forward pass of self-attention.
  
      Args:
          queries (torch.tensor): B x T1 x C1 tensor
          keys (torch.tensor): B x T2 x C2 tensor
          values (torch.tensor): B x T2 x C2 tensor
          q_mask (torch.tensor): B x T1 tensor, bool mask for variable length entries
          kv_mask (torch.tensor): B x T2 tensor, bool mask for variable length entries
      Output:
          attn_out (torch.tensor): B x T1 x C1 tensor
      """
      pos_q_seq = torch.arange(
          queries.size(-2), device=queries.device
      ).to(queries.dtype)
      pos_kv_seq = torch.arange(
          keys.size(-2), device=queries.device
      ).to(queries.dtype)
  
      pos_q_emb = self.text_pos_emb(pos_q_seq)
      pos_kv_emb = self.lm_pos_emb(pos_kv_seq)
  
      if q_mask is not None:
          pos_q_emb = pos_q_emb * q_mask.unsqueeze(2)
  
      if kv_mask is not None:
          pos_kv_emb = pos_kv_emb * kv_mask.unsqueeze(2)
  
      queries = (queries + pos_q_emb).transpose(1, 2)
      keys = (keys + pos_kv_emb).transpose(1, 2)
      values = (values + pos_kv_emb).transpose(1, 2)
  
      queries_enc = self.query_proj(queries).transpose(
          -2, -1
      )  # B x T1 x C1
      keys_enc = self.key_proj(keys)  # B x C1 x T2
      values_enc = self.value_proj(values).transpose(
          -2, -1
      )  # B x T2 x C1
  
      scores = (
          torch.matmul(queries_enc, keys_enc) / self.scale
      )  # B x T1 x T2
  
      if kv_mask is not None:
          scores.masked_fill_(~kv_mask.unsqueeze(-2), -float("inf"))
  
      return torch.matmul(torch.softmax(scores, dim=-1), values_enc)
  
imports: import math
  import torch
  from torch import nn
  from nemo.collections.tts.modules.submodules import ConvNorm
  from nemo.collections.tts.modules.transformer import PositionalEmbedding
label_desc: Use image node_id 2 for calling the SelfAttentionModule obj's underlying member method code with example usage: obj.forward(queries, keys, values, q_mask, kv_mask) and returns: torch