class_docstr: Parallel self-attention layer abstract class.
  
  Self-attention layer takes input with size [s, b, h]
  and returns output of the same size.
class_docstr_tok: ["Parallel","self-attention","layer","abstract","class",".","Self-attention","layer","takes","input","with","size","[","s",",","b",",","h","]","and","returns","output","of","the","same","size","."]
docstr: Forward method with activation checkpointing.
docstr_tok: ["Forward","method","with","activation","checkpointing","."]
code: def _checkpointed_attention_forward(
      self,
      query_layer,
      key_layer,
      value_layer,
      attention_mask,
      rotary_pos_emb=None,
      relative_position_bias=None,
      headscale_tensor=None,
      inference_mode=None,
  ):
      """Forward method with activation checkpointing."""
  
      def custom_forward(*inputs):
          if len(inputs) == 7:
              query_layer = inputs[0]
              key_layer = inputs[1]
              value_layer = inputs[2]
              attention_mask = inputs[3]
              rotary_pos_emb = inputs[4]
              relative_position_bias = inputs[5]
              headscale_tensor = inputs[6]
          elif len(inputs) == 8:
              query_layer = inputs[0]
              key_layer = inputs[1]
              value_layer = inputs[2]
              attention_mask = inputs[3]
              rotary_pos_emb = (inputs[4], inputs[5])
              relative_position_bias = inputs[6]
              headscale_tensor = inputs[7]
          else:
              raise ValueError("unexpected number of inputs")
          output_ = self.core_attention(
              query_layer,
              key_layer,
              value_layer,
              attention_mask,
              rotary_pos_emb=rotary_pos_emb,
              relative_position_bias=relative_position_bias,
              headscale_tensor=headscale_tensor,
              inference_mode=inference_mode,
          )
          return output_
  
      if rotary_pos_emb is None:
          rot_tuple = (rotary_pos_emb,)
      else:
          rot_tuple = (rotary_pos_emb[0], rotary_pos_emb[1])
  
      hidden_states = tensor_parallel.checkpoint(
          custom_forward,
          False,
          query_layer,
          key_layer,
          value_layer,
          attention_mask,
          *rot_tuple,
          relative_position_bias,
          headscale_tensor,
      )
  
      return hidden_states
  
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 ParallelAttention obj's underlying member method code with example usage: obj._checkpointed_attention_forward(query_layer, key_layer, value_layer, attention_mask, rotary_pos_emb, relative_position_bias, headscale_tensor, inference_mode) and returns: hidden_states, output_