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 the attention output of global tokens attending to all.
  
  Args:
      query (torch.Tensor): (batch, head, time, head_dim) The queries for global attention.
      key (torch.Tensor): (batch, head, time, head_dim) The keys for global attention.
      value (torch.Tensor): (batch, head, time, head_dim) The values for global attention.
      max_num_global_attn_indices (int): Maximum number of global attention indices in the batch.
      is_local_index_global_attn_nonzero (tuple): Non-padding values within global attention indices.
      is_index_global_attn_nonzero (tuple): Indices of global attention (non-zero elements).
      is_local_index_no_global_attn_nonzero (tuple): Padding values within global attention indices.
      is_index_masked (torch.Tensor): (batch, time) A boolean tensor indicating if an index is masked.
  
  Returns:
      global_attn_output (torch.Tensor): (batch, max_num_global_attn_indices, head x head_dim)
      The attention output of global tokens attending to all.
docstr_tok: ["Compute","the","attention","output","of","global","tokens","attending","to","all",".","Args",":","query","(","torch.Tensor",")",":","(","batch",",","head",",","time",",","head_dim",")","The","queries","for","global","attention",".","key","(","torch.Tensor",")",":","(","batch",",","head",",","time",",","head_dim",")","The","keys","for","global","attention",".","value","(","torch.Tensor",")",":","(","batch",",","head",",","time",",","head_dim",")","The","values","for","global","attention",".","max_num_global_attn_indices","(","int",")",":","Maximum","number","of","global","attention","indices","in","the","batch",".","is_local_index_global_attn_nonzero","(","tuple",")",":","Non-padding","values","within","global","attention","indices",".","is_index_global_attn_nonzero","(","tuple",")",":","Indices","of","global","attention","(","non-zero","elements",")",".","is_local_index_no_global_attn_nonzero","(","tuple",")",":","Padding","values","within","global","attention","indices",".","is_index_masked","(","torch.Tensor",")",":","(","batch",",","time",")","A","boolean","tensor","indicating","if","an","index","is","masked",".","Returns",":","global_attn_output","(","torch.Tensor",")",":","(","batch",",","max_num_global_attn_indices",",","head","x","head_dim",")","The","attention","output","of","global","tokens","attending","to","all","."]
code: def _compute_out_global_to_all(
      self,
      query: torch.Tensor,
      key: torch.Tensor,
      value: torch.Tensor,
      max_num_global_attn_indices: int,
      is_local_index_global_attn_nonzero: tuple,
      is_index_global_attn_nonzero: tuple,
      is_local_index_no_global_attn_nonzero: tuple,
      is_index_masked: torch.Tensor,
  ) -> Tuple[torch.Tensor, torch.Tensor]:
      """
      Compute the attention output of global tokens attending to all.
  
      Args:
          query (torch.Tensor): (batch, head, time, head_dim) The queries for global attention.
          key (torch.Tensor): (batch, head, time, head_dim) The keys for global attention.
          value (torch.Tensor): (batch, head, time, head_dim) The values for global attention.
          max_num_global_attn_indices (int): Maximum number of global attention indices in the batch.
          is_local_index_global_attn_nonzero (tuple): Non-padding values within global attention indices.
          is_index_global_attn_nonzero (tuple): Indices of global attention (non-zero elements).
          is_local_index_no_global_attn_nonzero (tuple): Padding values within global attention indices.
          is_index_masked (torch.Tensor): (batch, time) A boolean tensor indicating if an index is masked.
  
      Returns:
          global_attn_output (torch.Tensor): (batch, max_num_global_attn_indices, head x head_dim)
          The attention output of global tokens attending to all.
      """
  
      batch_size = key.shape[0]
      seq_len = key.shape[2]
  
      global_k = key.reshape(batch_size * self.h, -1, self.d_k)
      global_v = value.reshape(batch_size * self.h, -1, self.d_k)
  
      global_q = query.transpose(1, 2)
      global_q_from_global = global_q.new_zeros(
          batch_size, max_num_global_attn_indices, self.h, self.d_k
      )
      global_q_from_global[
          is_local_index_global_attn_nonzero
      ] = global_q[is_index_global_attn_nonzero]
      global_q_from_global = global_q_from_global.transpose(
          0, 1
      ).reshape(batch_size * self.h, -1, self.d_k)
  
      # compute attn scores
      global_attn_scores = torch.bmm(
          global_q_from_global, global_k.transpose(1, 2)
      )
      global_attn_scores = global_attn_scores.view(
          batch_size, self.h, max_num_global_attn_indices, seq_len
      )
  
      # need to transpose since ONNX export only supports consecutive indexing: https://pytorch.org/docs/stable/onnx.html#writes-sets
      global_attn_scores = global_attn_scores.transpose(1, 2)
      global_attn_scores[
          is_local_index_no_global_attn_nonzero[0],
          is_local_index_no_global_attn_nonzero[1],
          :,
          :,
      ] = torch.finfo(global_attn_scores.dtype).min
      global_attn_scores = global_attn_scores.transpose(1, 2)
  
      global_attn_scores = global_attn_scores.masked_fill(
          is_index_masked.transpose(2, 3),
          torch.finfo(global_attn_scores.dtype).min,
      )
  
      global_attn_scores = global_attn_scores.view(
          batch_size * self.h, max_num_global_attn_indices, seq_len
      )
  
      # compute global attn probs
      global_attn_probs_float = nn.functional.softmax(
          global_attn_scores, dim=-1, dtype=torch.float32
      )
  
      global_attn_probs = self.dropout(global_attn_probs_float)
  
      # global attn output
      global_attn_output = torch.bmm(global_attn_probs, global_v)
      global_attn_output = global_attn_output.view(
          batch_size, self.h, max_num_global_attn_indices, self.d_k
      )
  
      global_attn_output = global_attn_output[
          is_local_index_global_attn_nonzero[0],
          :,
          is_local_index_global_attn_nonzero[1],
      ]
  
      global_attn_output = global_attn_output.reshape(
          global_attn_output.shape[0], -1
      )
  
      return global_attn_output
  
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 6 for calling the RelPositionMultiHeadAttentionLongformer obj's underlying member method code with example usage: obj._compute_out_global_to_all(query, key, value, max_num_global_attn_indices, is_local_index_global_attn_nonzero, is_index_global_attn_nonzero, is_local_index_no_global_attn_nonzero, is_index_masked) and returns: global_attn_output