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: Compute Scaled Dot Product Local Attention with rel. positional encoding. using overlapping chunks
  Args:
      query (torch.Tensor): (batch, time, size)
      key (torch.Tensor): (batch, time, size)
      value(torch.Tensor): (batch, time, size)
      pad_mask (torch.Tensor): (batch, time)
      pos_emb (torch.Tensor) : (batch, 2w + 1, size)
      cache (torch.Tensor) : (batch, time_cache, size)
  Returns:
      output (torch.Tensor): transformed `value` (batch, time1, d_model) weighted by the query dot key attention
      cache (torch.Tensor) : (batch, time_cache_next, size)
docstr_tok: ["Compute","Scaled","Dot","Product","Local","Attention","with","rel",".","positional","encoding",".","using","overlapping","chunks","Args",":","query","(","torch.Tensor",")",":","(","batch",",","time",",","size",")","key","(","torch.Tensor",")",":","(","batch",",","time",",","size",")","value","(","torch.Tensor",")",":","(","batch",",","time",",","size",")","pad_mask","(","torch.Tensor",")",":","(","batch",",","time",")","pos_emb","(","torch.Tensor",")",":","(","batch",",","2w","+","1",",","size",")","cache","(","torch.Tensor",")",":","(","batch",",","time_cache",",","size",")","Returns",":","output","(","torch.Tensor",")",":","transformed","`","value","`","(","batch",",","time1",",","d_model",")","weighted","by","the","query","dot","key","attention","cache","(","torch.Tensor",")",":","(","batch",",","time_cache_next",",","size",")"]
code: def forward(self, query, key, value, pad_mask, pos_emb, cache=None):
      """Compute Scaled Dot Product Local Attention with rel. positional encoding. using overlapping chunks
      Args:
          query (torch.Tensor): (batch, time, size)
          key (torch.Tensor): (batch, time, size)
          value(torch.Tensor): (batch, time, size)
          pad_mask (torch.Tensor): (batch, time)
          pos_emb (torch.Tensor) : (batch, 2w + 1, size)
          cache (torch.Tensor) : (batch, time_cache, size)
      Returns:
          output (torch.Tensor): transformed `value` (batch, time1, d_model) weighted by the query dot key attention
          cache (torch.Tensor) : (batch, time_cache_next, size)
      """
  
      key, value, query, cache = self.update_cache(
          key=key, value=value, query=query, cache=cache
      )
  
      if torch.is_autocast_enabled():
          query, key, value = (
              query.to(torch.float32),
              key.to(torch.float32),
              value.to(torch.float32),
          )
  
      # temporary until we solve this more gracefully
      with avoid_float16_autocast_context():
          q, k, v = self.forward_qkv(query, key, value)
          n_batch, _, T, _ = q.size()
  
          w = max(self.att_context_size[0], self.att_context_size[1])
          if w <= 0:
              raise ValueError(
                  "When using local attention, context size must be set > 0"
              )
          pad_len = (2 * w - T % (2 * w)) % (2 * w)  # pad time to 2w
          q = F.pad(q, (0, 0, 0, pad_len))  # (batch, head, time, size)
          k = F.pad(k, (0, 0, 0, pad_len))  # (batch, head, time, size)
          v = F.pad(v, (0, 0, 0, pad_len))  # (batch, head, time, size)
          mask = F.pad(pad_mask, (0, pad_len), value=1.0)
  
          q_with_bias_u = q + self.pos_bias_u.unsqueeze(
              1
          )  # (batch, head, time, size)
          q_with_bias_v = q + self.pos_bias_v.unsqueeze(
              1
          )  # (batch, head, time, size)
  
          diagonal_matrix_ac = self.sliding_chunks_matmul_qk(
              q_with_bias_u, k, w, padding_value=0.0
          )  # (batch, head, time, 2w + 1)
  
          # add relative positional embedding
  
          n_batch_pos = pos_emb.size(0)
          p = (
              self.linear_pos(pos_emb)
              .view(n_batch_pos, -1, self.h, self.d_k)
              .transpose(1, 2)
          )
          # (batch, head, 2w, size)
          diagonal_matrix_bd = torch.matmul(
              q_with_bias_v, p.transpose(-2, -1)
          )
          # (batch, head, time, 2w + 1)
  
          start_pos = w - self.att_context_size[0]
          end_pos = w + self.att_context_size[1]
  
          diagonal_matrix_ac[
              :, :, :, : self.att_context_size[0]
          ] += diagonal_matrix_bd[:, :, :, : self.att_context_size[0]]
          diagonal_matrix_ac[
              :, :, :, -(self.att_context_size[1] + 1) :
          ] += diagonal_matrix_bd[:, :, :, self.att_context_size[0] :]
          scores = diagonal_matrix_ac / self.s_d_k
          # (batch, head, time, 2w + 1)
  
          # mask invalid positions
          scores[:, :, :, :start_pos] = -10000.0
          scores[:, :, :, end_pos + 1 :] = -10000.0
  
          # This implementation is fast and takes very little memory because num_heads x hidden_size = 1
          # from (bsz x seq_len) to (bsz x num_heads x seqlen x hidden_size)
          mask = mask.unsqueeze(dim=1).unsqueeze(dim=-1)
          # cast to float/half then replace 1's with -inf
          float_mask = mask.type_as(scores).masked_fill(mask, -10000.0)
          ones = float_mask.new_ones(
              size=float_mask.size()
          )  # tensor of ones
          # diagonal mask with zeros everywhere and -inf inplace of padding
          d_mask = self.sliding_chunks_matmul_qk(
              ones, float_mask, w, padding_value=0.0
          )
          # (batch, head, time, 2w + 1)
  
          scores += d_mask
  
          if self.global_tokens > 0:
              # create q, k, v for global attn
              if self.global_attn_separate:
                  global_q = self.global_q(query).view(
                      n_batch, -1, self.h, self.d_k
                  )
                  global_k = self.global_k(key).view(
                      n_batch, -1, self.h, self.d_k
                  )
                  global_v = self.global_v(value).view(
                      n_batch, -1, self.h, self.d_k
                  )
                  global_q = global_q.transpose(1, 2)
                  global_k = global_k.transpose(1, 2)
                  global_v = global_v.transpose(1, 2)
                  global_q = F.pad(
                      global_q, (0, 0, 0, pad_len)
                  )  # (batch, head, time, size)
                  global_k = F.pad(
                      global_k, (0, 0, 0, pad_len)
                  )  # (batch, head, time, size)
                  global_v = F.pad(
                      global_v, (0, 0, 0, pad_len)
                  )  # (batch, head, time, size)
              else:
                  global_q, global_k, global_v = q, k, v
  
              global_q /= self.s_d_k
  
              # assign which tokens are global
              is_index_global_attn = torch.zeros_like(pad_mask)
              is_index_global_attn[
                  :,
                  : self.global_tokens
                  * self.global_tokens_spacing : self.global_tokens_spacing,
              ] = 1.0
  
              # compute global attn indices
              (
                  max_num_global_attn_indices,
                  is_index_global_attn_nonzero,
                  is_local_index_global_attn_nonzero,
                  is_local_index_no_global_attn_nonzero,
              ) = self._get_global_attn_indices(
                  is_index_global_attn=is_index_global_attn
              )
  
              # calculate global attn probs with global keys
              # (batch, time, head, max_num_global_attn_indices)
              global_key_attn = self._compute_global_key_attn(
                  query=global_q.transpose(1, 2),
                  key=global_k.transpose(1, 2),
                  max_num_global_attn_indices=max_num_global_attn_indices,
                  is_index_global_attn_nonzero=is_index_global_attn_nonzero,
                  is_local_index_global_attn_nonzero=is_local_index_global_attn_nonzero,
                  is_local_index_no_global_attn_nonzero=is_local_index_no_global_attn_nonzero,
              ).transpose(1, 2)
  
              # concat to local_attn_probs
              # (batch, time, head, max_num_global_attn_indices + 2*w)
              scores = torch.cat((global_key_attn, scores), dim=-1)
  
              # free memory
              del global_key_attn
  
          attn = torch.softmax(scores, dim=-1).masked_fill(mask, 0.0)
          p_attn = self.dropout(attn)
          # (batch, head, time, 2w + 1)
  
          if self.global_tokens > 0:
              # compute sum of global and local attn
              out = self._compute_attn_output_with_global_indices(
                  value=v,
                  attn_probs=p_attn,
                  max_num_global_attn_indices=max_num_global_attn_indices,
                  is_index_global_attn_nonzero=is_index_global_attn_nonzero,
                  is_local_index_global_attn_nonzero=is_local_index_global_attn_nonzero,
                  w=w,
              )
          else:
              # compute local attn only
              out = self.sliding_chunks_matmul_pv(p_attn, v, w)
  
          out = out.reshape(n_batch, -1, self.h * self.d_k)[:, :T]
  
          if self.global_tokens > 0:
              out_global_to_all = self._compute_out_global_to_all(
                  query=global_q,
                  key=global_k,
                  value=global_v,
                  max_num_global_attn_indices=max_num_global_attn_indices,
                  is_local_index_global_attn_nonzero=is_local_index_global_attn_nonzero,
                  is_index_global_attn_nonzero=is_index_global_attn_nonzero,
                  is_local_index_no_global_attn_nonzero=is_local_index_no_global_attn_nonzero,
                  is_index_masked=mask,
              )
  
              # overwrite values with global attention
              out[is_index_global_attn_nonzero] = out_global_to_all
  
      ret = self.linear_out(out)
  
      if cache is None:
          return ret
      else:
          return ret, cache
  
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 2 for calling the RelPositionMultiHeadAttentionLongformer obj's underlying member method code with example usage: obj.forward(query, key, value, pad_mask, pos_emb, cache) and returns: ret, ret,  cache