class_docstr: Region where selective activation recomputation is applied.
  See Figure 3. in Reducing Activation Recomputation in Large Transformer Models
  https://arxiv.org/pdf/2205.05198.pdf for more details.
class_docstr_tok: ["Region","where","selective","activation","recomputation","is","applied",".","See","Figure","3.","in","Reducing","Activation","Recomputation","in","Large","Transformer","Models","https",":","\/\/arxiv.org\/pdf\/2205.05198.pdf","for","more","details","."]
code: def torch_attention(
      self,
      query_layer,
      key_layer,
      value_layer,
      attention_mask,
      attention_bias,
      inference_mode,
  ):
      sq, b, np, hn = query_layer.shape
      sk = key_layer.shape[0]
  
      if self.multi_query_attention:
          query_layer = rearrange(
              query_layer, "sq b np hn -> b (np sq) hn"
          )
          key_layer = rearrange(key_layer, "sk b 1 hn -> b hn sk")
          value_layer = rearrange(
              value_layer, "sv b np hn -> (b np) sv hn"
          )
      else:
          query_layer = rearrange(
              query_layer, "sq b np hn -> (b np) sq hn"
          )
          key_layer = rearrange(key_layer, "sk b np hn -> (b np) hn sk")
          value_layer = rearrange(
              value_layer, "sv b np hn -> (b np) sv hn"
          )
  
      matmul_input_buffer = torch.empty(
          query_layer.shape[0],
          query_layer.shape[1],
          key_layer.shape[2],
          dtype=query_layer.dtype,
          device=query_layer.device,
      )
  
      matmul_result = torch.baddbmm(
          matmul_input_buffer,
          query_layer,
          key_layer,
          beta=0.0,
          alpha=(1.0 / self.norm_factor)
          if self.normalize_attention_scores
          else 1.0,
      )
  
      # change view to [b, np, sq, sk]
      attention_scores = matmul_result.view(b, np, sq, sk)
  
      if attention_bias is not None:
          attention_scores += attention_bias
  
      attention_probs = self.scale_mask_softmax(
          attention_scores, attention_mask
      )
      # This is actually dropping out entire tokens to attend to, which might
      # seem a bit unusual, but is taken from the original Transformer paper.
  
      if not self.sequence_parallel:
          with tensor_parallel.random.get_cuda_rng_tracker().fork():
              attention_probs = self.attention_dropout(attention_probs)
      else:
          attention_probs = self.attention_dropout(attention_probs)
  
      # change view [b * np, sq, sk]
      attention_probs = rearrange(
          attention_probs, "b np sq sk -> (b np) sq sk"
      )
  
      # matmul: [b * np, sq, hn]
      context_layer = torch.bmm(attention_probs, value_layer)
  
      # change view [b, np, sq, hn]
      context_layer = rearrange(
          context_layer, "(b np) sq hn -> b np sq hn", np=np
      )
  
      return context_layer
  
imports: import math
  import torch
  import torch.nn.functional
  from einops import rearrange, repeat
  from nemo.collections.nlp.modules.common.megatron.adapters.parallel_adapters import AdapterName, InfusedAdapterConfig, LoraKQVAdapterConfig, LoraKQVAdapterWeightTyingConfig, LoraKVAdapterConfig, LoraQAdapterConfig
  from nemo.collections.nlp.modules.common.megatron.fused_softmax import MatchedScaleMaskSoftmax
  from nemo.collections.nlp.modules.common.megatron.module import MegatronModule
  from nemo.collections.nlp.modules.common.megatron.position_embedding import XPOSPositionEmbedding
  from nemo.collections.nlp.modules.common.megatron.position_embedding.rotary_position_embedding import apply_rotary_pos_emb
  from nemo.collections.nlp.modules.common.megatron.utils import ApexGuardDefaults, _cast_if_autocast_enabled, attention_mask_func
  from nemo.core import adapter_mixins
label_desc: Use image node_id 3 for calling the CoreAttention obj's underlying member method code with example usage: obj.torch_attention(query_layer, key_layer, value_layer, attention_mask, attention_bias, inference_mode) and returns: context_layer