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 forward(
      self,
      query_layer,
      key_layer,
      value_layer,
      attention_mask,
      layer_past=None,
      get_key_value=False,
      rotary_pos_emb=None,
      relative_position_bias=None,
      headscale_tensor=None,
      inference_mode=None,
  ):
      b, np, sq, sk, hn = (
          query_layer.size(1),
          query_layer.size(2),
          query_layer.size(0),
          key_layer.size(0),
          query_layer.size(3),
      )
  
      # ==================================================
      # Update attention mask for inference. [b, np, sq, sk]
      # ==================================================
      if get_key_value:
          with torch.no_grad():
              if layer_past is not None:
                  attention_mask = attention_mask[
                      ..., sq - 1, :sk
                  ].unsqueeze(2)
              else:
                  attention_mask = attention_mask[..., :sq, :sk]
  
      # ==================================================
      # Update attention bias. [b, np, sq, sk]
      # ==================================================
      if relative_position_bias is not None:
          relative_position_bias = relative_position_bias[
              :,
              self.num_attention_heads_partition_offset : self.num_attention_heads_partition_offset
              + self.num_attention_heads_per_partition,
              -sq:,
              -sk:,
          ]
  
      # ==================================================
      # Update query_layer, key_layer, value_layer
      # ==================================================
      # TODO: figure out how to do this
      # apply relative positional encoding (rotary embedding)
      if rotary_pos_emb is not None:
          q_pos_emb, k_pos_emb = rotary_pos_emb
          query_layer = apply_rotary_pos_emb(query_layer, q_pos_emb)
          key_layer = apply_rotary_pos_emb(key_layer, k_pos_emb)
          # TODO, can apply positional embedding to value_layer so it has
          # absolute positional embedding.
          # otherwise, only relative positional embedding takes effect
          # value_layer = apply_rotary_pos_emb(value_layer, k_pos_emb)
  
      if self.position_embedding_type.lower() == "xpos":
          query_layer = self.xpos(
              query_layer,
              offset=key_layer.shape[-2] - query_layer.shape[-2],
              downscale=False,
          )
          key_layer = self.xpos(key_layer, offset=0, downscale=True)
  
      # ==================================================
      # query_layer [sq, b, np, hn]
      # key_layer   [sk, b, np, hn]
      # value_layer [sk, b, np, hn]
      # attention_mask [b, 1, sq, sk] or [b, s]
      # relative_position_bias [b, np, sq, sk]
      # context_layer [b, np, sq, hn]
      # ==================================================
      context_layer = self.attn_fn(
          query_layer,
          key_layer,
          value_layer,
          attention_mask,
          relative_position_bias,
          inference_mode,
      )
  
      if headscale_tensor is not None:
          context_layer = context_layer * headscale_tensor
  
      # [b, np, sq, hn] --> [sq, b, np, hn]
      context_layer = context_layer.permute(2, 0, 1, 3).contiguous()
  
      # [sq, b, np, hn] --> [sq, b, hp]
      new_context_layer_shape = context_layer.size()[:-2] + (
          self.hidden_size_per_partition,
      )
      context_layer = context_layer.view(*new_context_layer_shape)
  
      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 2 for calling the CoreAttention obj's underlying member method code with example usage: obj.forward(query_layer, key_layer, value_layer, attention_mask, layer_past, get_key_value, rotary_pos_emb, relative_position_bias, headscale_tensor, inference_mode) and returns: context_layer