class_docstr: MEGA module.
  
  Args:
      size: Input/Output size.
      num_heads: Number of EMA heads.
      qk_size: Shared query and key size for attention module.
      v_size: Value size for attention module.
      qk_v_size: (QK, V) sizes for attention module.
      activation: Activation function type.
      normalization: Normalization module.
      rel_pos_bias_type: Type of relative position bias in attention module.
      max_positions: Maximum number of position for RelativePositionBias.
      truncation_length: Maximum length for truncation in EMA module.
      chunk_size: Chunk size for attention computation (-1 = full context).
      dropout_rate: Dropout rate for inner modules.
      att_dropout_rate: Dropout rate for the attention module.
      ema_dropout_rate: Dropout rate for the EMA module.
docstr: Compute attention weights with softmax.
  
  Args:
      query: Query tensor. (B, 1, L, D)
      key: Key tensor. (B, 1, L, D)
      mask: Sequence mask. (B, 1, L)
      attn_mask: Attention mask. (1, L, L)
  
  Returns:
      attn_weights: Attention weights. (B, 1, L, L)
docstr_tok: ["Compute","attention","weights","with","softmax",".","Args",":","query",":","Query","tensor",".","(","B",",","1",",","L",",","D",")","key",":","Key","tensor",".","(","B",",","1",",","L",",","D",")","mask",":","Sequence","mask",".","(","B",",","1",",","L",")","attn_mask",":","Attention","mask",".","(","1",",","L",",","L",")","Returns",":","attn_weights",":","Attention","weights",".","(","B",",","1",",","L",",","L",")"]
code: def softmax_attention(
      self,
      query: torch.Tensor,
      key: torch.Tensor,
      mask: Optional[torch.Tensor] = None,
      attn_mask: Optional[torch.Tensor] = None,
  ) -> torch.Tensor:
      """Compute attention weights with softmax.
  
      Args:
          query: Query tensor. (B, 1, L, D)
          key: Key tensor. (B, 1, L, D)
          mask: Sequence mask. (B, 1, L)
          attn_mask: Attention mask. (1, L, L)
  
      Returns:
          attn_weights: Attention weights. (B, 1, L, L)
  
      """
      length = key.size(2)
  
      bias = self.rel_pos_bias(length)
  
      if length != query.size(2):
          bias = bias[-1:]
  
      query = query * self.scaling
  
      qk = torch.matmul(query, key.transpose(2, 3)) + bias
  
      if attn_mask is not None:
          qk = qk.masked_fill(attn_mask.unsqueeze(1), float("-inf"))
  
      if mask is not None:
          mask_all = mask.all(dim=-1, keepdim=True)
          mask = torch.logical_and(mask, ~mask_all)
  
          qk = qk.masked_fill(mask.unsqueeze(2), float("-inf"))
  
      attn_weights = torch.softmax(
          qk, dim=-1, dtype=torch.float32
      ).type_as(qk)
  
      return attn_weights
  
file_path: espnet/espnet2/asr_transducer/decoder/blocks/mega.py
imports: from typing import Dict, Optional, Tuple
  import torch
  from espnet2.asr_transducer.decoder.modules.mega.multi_head_damped_ema import MultiHeadDampedEMA
  from espnet2.asr_transducer.decoder.modules.mega.positional_bias import RelativePositionBias, RotaryRelativePositionBias
label_desc: Use image node_id 3 for calling the MEGA obj's underlying member method code with example usage: obj.softmax_attention(query, key, mask, attn_mask) and returns: attn_weights