class_docstr: Multi-Head Attention layer of Transformer-XL with sliding window local+global attention from Longformer.
  Partially adapted from allenai (https://github.com/allenai/longformer/blob/master/longformer/sliding_chunks.py)
  and huggingface (https://github.com/huggingface/transformers/blob/main/src/transformers/models/longformer/modeling_longformer.py) 
  Paper: https://arxiv.org/abs/1901.02860 (Transformer-XL),
         https://arxiv.org/abs/2004.05150 (Longformer)
  Args:
      n_head (int): number of heads
      n_feat (int): size of the features
      dropout_rate (float): dropout rate
      pos_bias_u (Tensor): the positional bias matrix U
      pos_bias_v (Tensor): the positional bias matrix V
      att_context_size (List[int]): List of 2 ints corresponding to left and right attention context sizes.
      max_cache_len (int): the maximum size of cache
      global_tokens (int): number of tokens to be used for global attention
      global_tokens_spacing (int): how far apart the global tokens are
      global_attn_separate (bool): whether the q, k, v layers used for global tokens should be separate
class_docstr_tok: ["Multi-Head","Attention","layer","of","Transformer-XL","with","sliding","window","local+global","attention","from","Longformer",".","Partially","adapted","from","allenai","(","https",":","\/\/github.com\/allenai\/longformer\/blob\/master\/longformer\/sliding_chunks.py",")","and","huggingface","(","https",":","\/\/github.com\/huggingface\/transformers\/blob\/main\/src\/transformers\/models\/longformer\/modeling_longformer.py",")","Paper",":","https",":","\/\/arxiv.org\/abs\/1901.02860","(","Transformer-XL",")",",","https",":","\/\/arxiv.org\/abs\/2004.05150","(","Longformer",")","Args",":","n_head","(","int",")",":","number","of","heads","n_feat","(","int",")",":","size","of","the","features","dropout_rate","(","float",")",":","dropout","rate","pos_bias_u","(","Tensor",")",":","the","positional","bias","matrix","U","pos_bias_v","(","Tensor",")",":","the","positional","bias","matrix","V","att_context_size","(","List","[","int","]",")",":","List","of","2","ints","corresponding","to","left","and","right","attention","context","sizes",".","max_cache_len","(","int",")",":","the","maximum","size","of","cache","global_tokens","(","int",")",":","number","of","tokens","to","be","used","for","global","attention","global_tokens_spacing","(","int",")",":","how","far","apart","the","global","tokens","are","global_attn_separate","(","bool",")",":","whether","the","q",",","k",",","v","layers","used","for","global","tokens","should","be","separate"]
docstr: Matrix multiplication of query x key tensors using with a sliding window attention pattern.
  This implementation splits the input into overlapping chunks of size 2w
  with an overlap of size w
  
  Args:
      q (torch.Tensor): (batch, head, time, size)
      k (torch.Tensor): (batch, head, time, size)
      w (int): Chunk overlap size
      padding_value (float): Value to pad with
  
  Returns:
      output (torch.Tensor): (batch, head, time, 2w + 1)
docstr_tok: ["Matrix","multiplication","of","query","x","key","tensors","using","with","a","sliding","window","attention","pattern",".","This","implementation","splits","the","input","into","overlapping","chunks","of","size","2w","with","an","overlap","of","size","w","Args",":","q","(","torch.Tensor",")",":","(","batch",",","head",",","time",",","size",")","k","(","torch.Tensor",")",":","(","batch",",","head",",","time",",","size",")","w","(","int",")",":","Chunk","overlap","size","padding_value","(","float",")",":","Value","to","pad","with","Returns",":","output","(","torch.Tensor",")",":","(","batch",",","head",",","time",",","2w","+","1",")"]
code: def sliding_chunks_matmul_qk(
      self,
      q: torch.Tensor,
      k: torch.Tensor,
      w: int,
      padding_value: float,
  ) -> torch.Tensor:
      """Matrix multiplication of query x key tensors using with a sliding window attention pattern.
      This implementation splits the input into overlapping chunks of size 2w
      with an overlap of size w
  
      Args:
          q (torch.Tensor): (batch, head, time, size)
          k (torch.Tensor): (batch, head, time, size)
          w (int): Chunk overlap size
          padding_value (float): Value to pad with
  
      Returns:
          output (torch.Tensor): (batch, head, time, 2w + 1)
      """
      bsz, num_heads, seqlen, head_dim = q.size()
      assert seqlen % (w * 2) == 0
      assert q.size() == k.size()
  
      chunks_count = seqlen // w - 1
  
      # group bsz and num_heads dimensions into one, then chunk seqlen into chunks of size w * 2
      q = q.reshape(bsz * num_heads, seqlen, head_dim)
      k = k.reshape(bsz * num_heads, seqlen, head_dim)
  
      chunk_q = self._chunk_overlap(
          q, w
      )  # (batch x head, chunk_count, 2w, size)
      chunk_k = self._chunk_overlap(
          k, w
      )  # (batch x head, chunk_count, 2w, size)
  
      # matrix multipication
      # bcxd: bsz*num_heads x chunks x 2w x head_dim
      # bcyd: bsz*num_heads x chunks x 2w x head_dim
      # bcxy: bsz*num_heads x chunks x 2w x 2w
      chunk_attn = torch.einsum(
          "bcxd,bcyd->bcxy", (chunk_q, chunk_k)
      )  # multiply
      # (batch x head, chunk_count, 2w, 2w)
  
      # convert diagonals into columns
      diagonal_chunk_attn = self._skew(
          chunk_attn,
          direction=(0, 0, 0, 1),
          padding_value=padding_value,
      )
      # (batch x head, chunk_count, 2w, 2w + 1)
  
      # allocate space for the overall attention matrix where the chunks are combined. The last dimension
      # has (w * 2 + 1) columns. The first (w) columns are the w lower triangles (attention from a word to
      # w previous words). The following column is attention score from each word to itself, then
      # followed by w columns for the upper triangle.
  
      diagonal_attn = diagonal_chunk_attn.new_empty(
          (bsz * num_heads, chunks_count + 1, w, w * 2 + 1)
      )
      # (batch x head, chunk_count + 1, w, 2w + 1)
  
      # copy parts from diagonal_chunk_attn into the compined matrix of attentions
      # - copying the main diagonal and the upper triangle
      diagonal_attn[:, :-1, :, w:] = diagonal_chunk_attn[
          :, :, :w, : w + 1
      ]
      diagonal_attn[:, -1, :, w:] = diagonal_chunk_attn[
          :, -1, w:, : w + 1
      ]
      # - copying the lower triangle
      diagonal_attn[:, 1:, :, :w] = diagonal_chunk_attn[
          :, :, -(w + 1) : -1, w + 1 :
      ]
      diagonal_attn[:, 0, 1:w, 1:w] = diagonal_chunk_attn[
          :, 0, : w - 1, 1 - w :
      ]
  
      # separate bsz and num_heads dimensions again
      diagonal_attn = diagonal_attn.view(
          bsz, num_heads, seqlen, 2 * w + 1
      )
      # (batch, head, time, 2w + 1)
  
      self.mask_invalid_locations(diagonal_attn, w)
  
      return diagonal_attn
  
imports: import math
  from functools import lru_cache
  from typing import List, Tuple
  import torch
  import torch.nn
  import torch.nn.functional
  from nemo.utils import avoid_float16_autocast_context
label_desc: Use image node_id 12 for calling the RelPositionMultiHeadAttentionLongformer obj's underlying member method code with example usage: obj.sliding_chunks_matmul_qk(q, k, w, padding_value) and returns: diagonal_attn