class_docstr: Attention layer abstract class.
  
  This layer only contains common modules required for the "self attn" and
  "cross attn" specializations.
class_docstr_tok: ["Attention","layer","abstract","class",".","This","layer","only","contains","common","modules","required","for","the","``","self","attn","''","and","''","cross","attn","''","specializations","."]
docstr: Forward method with selective activation checkpointing.
docstr_tok: ["Forward","method","with","selective","activation","checkpointing","."]
code: def _checkpointed_attention_forward(
      self,
      query,
      key,
      value,
      attention_mask,
      rotary_pos_emb=None,
      attn_mask_type=None,
  ):
      """Forward method with selective activation checkpointing."""
  
      def custom_forward(*inputs):
          query = inputs[0]
          key = inputs[1]
          value = inputs[2]
          attention_mask = inputs[3]
          attn_mask_type = inputs[5]
          attn_mask_type = AttnMaskType(attn_mask_type.item())
          output_ = self.core_attention(
              query,
              key,
              value,
              attention_mask,
              attn_mask_type=attn_mask_type,
          )
          return output_
  
      if attn_mask_type is None:
          attn_mask_type = self.attn_mask_type
      attn_mask_type = torch.tensor(
          [attn_mask_type.value], dtype=torch.int
      )
      hidden_states = tensor_parallel.checkpoint(
          custom_forward,
          False,
          query,
          key,
          value,
          attention_mask,
          rotary_pos_emb,
          attn_mask_type,
      )
  
      return hidden_states
  
imports: from abc import ABC, abstractmethod
  from dataclasses import dataclass
  from typing import Union
  import torch
  from megatron.core import parallel_state, tensor_parallel
  from megatron.core.models.common.embeddings.rotary_pos_embedding import apply_rotary_pos_emb
  from megatron.core.transformer.custom_layers.transformer_engine import SplitAlongDim
  from megatron.core.transformer.enums import AttnMaskType
  from megatron.core.transformer.identity_op import IdentityFuncOp, IdentityOp
  from megatron.core.transformer.module import MegatronModule
  from megatron.core.transformer.spec_utils import ModuleSpec, build_module
  from megatron.core.transformer.transformer_config import TransformerConfig
  from megatron.core.utils import divide
  from .enums import AttnMaskType
  from .transformer_config import TransformerConfig
  from .utils import make_sharded_tensors_for_checkpoint
label_desc: Use image node_id 2 for calling the Attention obj's underlying member method code with example usage: obj._checkpointed_attention_forward(query, key, value, attention_mask, rotary_pos_emb, attn_mask_type) and returns: hidden_states, output_