class_docstr: Lunar Multi-headed attention.
  See "Linformer: Self-Attention with Linear Complexity" for more details.
docstr: Input shape: Time x Batch x Channel
  Args:
      context_padding_mask (ByteTensor, optional): mask to exclude
          keys that are pads, of shape `(batch, src_len)`, where
          padding elements are indicated by 1s.
      need_weights (bool, optional): return the attention weights,
          averaged over heads (default: False).
      need_head_weights (bool, optional): return the attention
          weights for each head. Implies *need_weights*. Default:
          return the average attention weights over all heads.
docstr_tok: ["Input","shape",":","Time","x","Batch","x","Channel","Args",":","context_padding_mask","(","ByteTensor",",","optional",")",":","mask","to","exclude","keys","that","are","pads",",","of","shape","`","(","batch",",","src_len",")","`",",","where","padding","elements","are","indicated","by","1s",".","need_weights","(","bool",",","optional",")",":","return","the","attention","weights",",","averaged","over","heads","(","default",":","False",")",".","need_head_weights","(","bool",",","optional",")",":","return","the","attention","weights","for","each","head",".","Implies","*","need_weights","*",".","Default",":","return","the","average","attention","weights","over","all","heads","."]
code: def forward(
      self,
      query,
      pquery,
      context: Optional[Tensor],
      context_padding_mask: Optional[Tensor] = None,
      pcontext_padding_mask: Optional[Tensor] = None,
      incremental_state: Optional[
          Dict[str, Dict[str, Optional[Tensor]]]
      ] = None,
      need_weights: bool = False,
      static_context: bool = False,
      need_head_weights: bool = False,
  ) -> Tuple[Tensor, Tensor, Optional[Tensor]]:
      """Input shape: Time x Batch x Channel
      Args:
          context_padding_mask (ByteTensor, optional): mask to exclude
              keys that are pads, of shape `(batch, src_len)`, where
              padding elements are indicated by 1s.
          need_weights (bool, optional): return the attention weights,
              averaged over heads (default: False).
          need_head_weights (bool, optional): return the attention
              weights for each head. Implies *need_weights*. Default:
              return the average attention weights over all heads.
      """
      if need_head_weights:
          need_weights = True
  
      tgt_len, bsz, embed_dim = query.size()
      assert embed_dim == self.embed_dim
      assert (
          not self.self_attention or incremental_state is None
      ), "For incremental self attention (causal attention), please use LunarCausalAttention"
  
      if self.self_attention:
          context = query
  
      if incremental_state is not None:
          saved_state = self._get_input_buffer(incremental_state)
          if saved_state is not None and "prev_key" in saved_state:
              # previous time steps are cached - no need to recompute
              # key and value if they are static
              if static_context:
                  assert (
                      self.encoder_decoder_attention
                      and not self.self_attention
                  )
                  context = None
      else:
          saved_state = None
  
      # L x B x D
      pcontext = self.compute_pcontext(
          query,
          pquery,
          context,
          context_padding_mask,
          incremental_state,
          static_context,
      )
  
      key_padding_mask = pcontext_padding_mask
  
      q = self.q_proj(query)
      if pcontext is None:
          assert context is None
          k = v = None
      elif self.c_proj is not None:
          k = v = (
              self.c_proj(pcontext)
              .view(-1, bsz * self.num_heads, self.head_dim)
              .transpose(0, 1)
          )
      else:
          k = (
              self.k_proj(pcontext)
              .view(-1, bsz * self.num_heads, self.head_dim)
              .transpose(0, 1)
          )
          v = (
              self.v_proj(pcontext)
              .view(-1, bsz * self.num_heads, self.head_dim)
              .transpose(0, 1)
          )
  
      q = q * self.scaling
      q = (
          q.contiguous()
          .view(tgt_len, bsz * self.num_heads, self.head_dim)
          .transpose(0, 1)
      )
  
      if saved_state is not None:
          # saved states are stored with shape (bsz, num_heads, seq_len, head_dim)
          if "prev_key" in saved_state:
              _prev_key = saved_state["prev_key"]
              assert _prev_key is not None
              prev_key = _prev_key.view(
                  bsz * self.num_heads, -1, self.head_dim
              )
              if static_context:
                  k = prev_key
              else:
                  assert k is not None
                  k = torch.cat([prev_key, k], dim=1)
          if "prev_value" in saved_state:
              _prev_value = saved_state["prev_value"]
              assert _prev_value is not None
              prev_value = _prev_value.view(
                  bsz * self.num_heads, -1, self.head_dim
              )
              if static_context:
                  v = prev_value
              else:
                  assert v is not None
                  v = torch.cat([prev_value, v], dim=1)
          # pcontext are stored with shape (bsz, proj_len, model_dim)
          if "prev_pcontext" in saved_state:
              # TODO save prev_pcontext for causal attention
              _prev_pcontext = saved_state["prev_pcontext"]
              assert _prev_pcontext is not None
              prev_pcontext = _prev_pcontext.transpose(0, 1)
              if static_context:
                  pcontext = prev_pcontext
              else:
                  raise RuntimeError("pcontext error")
  
          assert (
              k is not None and v is not None and pcontext is not None
          )
          saved_state["prev_key"] = k.view(
              bsz, self.num_heads, -1, self.head_dim
          )
          saved_state["prev_value"] = v.view(
              bsz, self.num_heads, -1, self.head_dim
          )
          saved_state["prev_pcontext"] = pcontext.transpose(0, 1)
          # In this branch incremental_state is never None
          assert incremental_state is not None
          incremental_state = self._set_input_buffer(
              incremental_state, saved_state
          )
  
      assert k is not None and v is not None and pcontext is not None
      src_len = k.size(1)
  
      # This is part of a workaround to get around fork/join parallelism
      # not supporting Optional types.
      if key_padding_mask is not None and key_padding_mask.dim() == 0:
          key_padding_mask = None
  
      if key_padding_mask is not None:
          assert key_padding_mask.size(0) == bsz
          assert key_padding_mask.size(1) == src_len
  
      attn_weights = torch.bmm(q, k.transpose(1, 2))
      attn_weights = LunarMultiheadAttention.apply_sparse_mask(
          attn_weights, tgt_len, src_len, bsz
      )
  
      assert list(attn_weights.size()) == [
          bsz * self.num_heads,
          tgt_len,
          src_len,
      ]
  
      if key_padding_mask is not None:
          # don't attend to padding symbols
          attn_weights = attn_weights.view(
              bsz, self.num_heads, tgt_len, src_len
          )
          if not self.tpu:
              attn_weights = attn_weights.masked_fill(
                  key_padding_mask.unsqueeze(1)
                  .unsqueeze(2)
                  .to(torch.bool),
                  float("-inf"),
              )
          else:
              attn_weights = attn_weights.transpose(0, 2)
              attn_weights = attn_weights.masked_fill(
                  key_padding_mask, float("-inf")
              )
              attn_weights = attn_weights.transpose(0, 2)
          attn_weights = attn_weights.view(
              bsz * self.num_heads, tgt_len, src_len
          )
  
      attn_weights_float = utils.softmax(
          attn_weights, dim=-1, onnx_trace=self.onnx_trace
      )
      attn_weights = attn_weights_float.type_as(attn_weights)
      attn_probs = self.dropout_module(attn_weights)
  
      attn = torch.bmm(attn_probs, v)
      assert list(attn.size()) == [
          bsz * self.num_heads,
          tgt_len,
          self.head_dim,
      ]
      if self.onnx_trace and attn.size(1) == 1:
          # when ONNX tracing a single decoder step (sequence length == 1)
          # the transpose is a no-op copy before view, thus unnecessary
          attn = attn.contiguous().view(tgt_len, bsz, embed_dim)
      else:
          attn = (
              attn.transpose(0, 1)
              .contiguous()
              .view(tgt_len, bsz, embed_dim)
          )
      attn = self.out_proj(attn)
      attn_weights: Optional[Tensor] = None
      if need_weights:
          attn_weights = attn_weights_float.view(
              bsz, self.num_heads, tgt_len, src_len
          ).transpose(1, 0)
          if not need_head_weights:
              # average attention weights over heads
              attn_weights = attn_weights.mean(dim=0)
  
      return attn, pcontext, attn_weights
  
imports: import math
  from typing import Dict, Optional, Tuple, Union
  import torch
  import torch.nn.functional
  from fairseq import utils
  from fairseq.incremental_decoding_utils import with_incremental_state
  from fairseq.modules.fairseq_dropout import FairseqDropout
  from torch import nn, Tensor
label_desc: Use image node_id 7 for calling the LunarMultiheadAttention obj's underlying member method code with example usage: obj.forward(query, pquery, context, context_padding_mask, pcontext_padding_mask, incremental_state, need_weights, static_context, need_head_weights) and returns: attn,  pcontext,  attn_weights