class_docstr: Lunar Causal attention.
  See "Linformer: Self-Attention with Linear Complexity" for more details.
docstr: Input shape: Time x Batch x Channel
  Args:
      key_padding_mask (ByteTensor, optional): mask to exclude
          keys that are pads, of shape `(batch, src_len)`, where
          padding elements are indicated by 1s.
      pkey_padding_mask (ByteTensor, optional): mask to exclude
          keys that are pads, of shape `(batch, proj_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",":","key_padding_mask","(","ByteTensor",",","optional",")",":","mask","to","exclude","keys","that","are","pads",",","of","shape","`","(","batch",",","src_len",")","`",",","where","padding","elements","are","indicated","by","1s",".","pkey_padding_mask","(","ByteTensor",",","optional",")",":","mask","to","exclude","keys","that","are","pads",",","of","shape","`","(","batch",",","proj_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,
      key_padding_mask: Optional[Tensor] = None,
      pkey_padding_mask: Optional[Tensor] = None,
      incremental_state: Optional[
          Dict[str, Dict[str, Optional[Tensor]]]
      ] = None,
      need_weights: bool = False,
      need_head_weights: bool = False,
  ) -> Tuple[Tensor, Optional[Tensor]]:
      """Input shape: Time x Batch x Channel
      Args:
          key_padding_mask (ByteTensor, optional): mask to exclude
              keys that are pads, of shape `(batch, src_len)`, where
              padding elements are indicated by 1s.
          pkey_padding_mask (ByteTensor, optional): mask to exclude
              keys that are pads, of shape `(batch, proj_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
  
      pq = None
      num_steps = None
      if incremental_state is not None:
          saved_state = self._get_input_buffer(incremental_state)
          if saved_state is not None and "prev_pquery" in saved_state:
              # previous time steps are cached - no need to recompute pquery
              # B x H x L x K
              pq = saved_state["prev_pquery"]
          key_accum_mat = 0
          value_accum_mat = 0
      else:
          saved_state = None
          key_accum_mat = None
          value_accum_mat = None
  
      if pq is None:
          plen = pquery.size(0)
          # L x B x D -> L x B x H x K
          pq = self.pq_proj(pquery).view(
              plen, bsz, self.num_heads, self.head_dim
          )
          # L x B x H x K -> B x H x L x K
          pq = pq.permute(1, 2, 0, 3) * self.scaling
  
      plen = pq.size(2)
      # B*H x N x L
      pattn_weights = self._compute_pattention(
          pq, query, key_padding_mask
      )
      pattn_weights = self.dropout_module(pattn_weights)
  
      # N x B x D -> B*H x N x K
      q = self.q_proj(query) * self.scaling
      q = q.view(
          tgt_len, bsz * self.num_heads, self.head_dim
      ).transpose(0, 1)
      # N x B x D -> B*H x N x K
      if self.c_proj is not None:
          k = v = (
              self.c_proj(query)
              .view(tgt_len, bsz * self.num_heads, self.head_dim)
              .transpose(0, 1)
          )
      else:
          k = (
              self.k_proj(query)
              .view(tgt_len, bsz * self.num_heads, self.head_dim)
              .transpose(0, 1)
          )
          v = (
              self.v_proj(query)
              .view(tgt_len, bsz * self.num_heads, self.head_dim)
              .transpose(0, 1)
          )
  
      efficient_causal_attention = (
          efficient_causal_attention_parallel
          if self.parallel
          else efficient_causal_attention_seq
      )
  
      if saved_state is not None:
          # key accumulative matrix are store with shape (bsz, num_heads, head_dim, plen)
          if "prev_key_accum_mat" in saved_state:
              _prev_key_accum_mat = saved_state["prev_key_accum_mat"]
              key_accum_mat = _prev_key_accum_mat.view(
                  bsz * self.num_heads, self.head_dim, plen
              )
          # value accumulative matrix are store with shape (bsz, num_heads, plen, head_dim)
          if "prev_value_accum_mat" in saved_state:
              _prev_value_accum_mat = saved_state[
                  "prev_value_accum_mat"
              ]
              value_accum_mat = _prev_value_accum_mat.view(
                  bsz * self.num_heads, plen, self.head_dim
              )
          if "prev_num_steps" in saved_state:
              _prev_num_steps = saved_state["prev_num_steps"]
              num_steps = (
                  _prev_num_steps.view(bsz * self.num_heads) + 1.0
              )
  
      if num_steps is None:
          num_steps = query.new_ones(bsz * self.num_heads)
  
      # This is part of a workaround to get around fork/join parallelism
      # not supporting Optional types.
      if pkey_padding_mask is not None and pkey_padding_mask.dim() == 0:
          pkey_padding_mask = None
  
      if incremental_state is not None:
          attn_weights, key_accum_mat = incremental_causal_attention(
              q, k, pattn_weights, key_accum_mat, num_steps
          )
      else:
          attn_weights = efficient_causal_attention(q, k, pattn_weights)
  
      assert list(attn_weights.size()) == [
          bsz * self.num_heads,
          tgt_len,
          plen,
      ]
  
      if pkey_padding_mask is not None:
          # don't attend to padding symbols
          attn_weights = attn_weights.view(
              bsz, self.num_heads, tgt_len, plen
          )
          if not self.tpu:
              attn_weights = attn_weights.masked_fill(
                  pkey_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(
                  pkey_padding_mask, float("-inf")
              )
              attn_weights = attn_weights.transpose(0, 2)
          attn_weights = attn_weights.view(
              bsz * self.num_heads, tgt_len, plen
          )
  
      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)
  
      if incremental_state is not None:
          attn, value_accum_mat = incremental_causal_attention(
              attn_probs, pattn_weights, v, value_accum_mat, num_steps
          )
      else:
          attn = efficient_causal_attention(
              attn_probs, pattn_weights, v
          )
  
      if saved_state is not None:
          saved_state["prev_pquery"] = pq
          saved_state["prev_key_accum_mat"] = key_accum_mat.view(
              bsz, self.num_heads, self.head_dim, plen
          )
          saved_state["prev_value_accum_mat"] = value_accum_mat.view(
              bsz, self.num_heads, plen, self.head_dim
          )
          saved_state["prev_num_steps"] = num_steps.view(
              bsz, self.num_heads
          )
          # 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 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, plen
          ).transpose(1, 0)
          if not need_head_weights:
              # average attention weights over heads
              attn_weights = attn_weights.mean(dim=0)
  
      return attn, 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 6 for calling the LunarCausalAttention obj's underlying member method code with example usage: obj.forward(query, pquery, key_padding_mask, pkey_padding_mask, incremental_state, need_weights, need_head_weights) and returns: attn,  attn_weights