class_docstr: Wrapper for the Transformer-Engine's `DotProductAttention` layer that also
  has "flash attention" enabled.
  
  Note that if Megatron's parallel_state has not been initialized yet, the
  tp_group and cp_group passed to TE will be None and must be set later
  via set_tensor_parallel_group() and set_context_parallel_group().
code: def forward(
      self,
      query: Tensor,
      key: Tensor,
      value: Tensor,
      attention_mask: Tensor,
      attn_mask_type: AttnMaskType,
  ):
      if self.te_forward_mask_type:
          return super().forward(
              query,
              key,
              value,
              attention_mask,
              attn_mask_type=attn_mask_type.name,
          )
      else:
          return super().forward(query, key, value, attention_mask)
  
imports: import os
  from importlib.metadata import version
  from typing import Callable
  import torch
  import transformer_engine
  from pkg_resources import packaging
  from torch import Tensor
  from megatron.core import ModelParallelConfig
  from megatron.core.parallel_state import get_context_parallel_global_ranks, get_context_parallel_group, get_tensor_model_parallel_group
  from megatron.core.tensor_parallel import get_cuda_rng_tracker
  from megatron.core.transformer.enums import AttnMaskType
  from megatron.core.transformer.transformer_config import TransformerConfig
  from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint
label_desc: Use image node_id 2 for calling the TEDotProductAttention obj's underlying member method code with example usage: obj.forward(query, key, value, attention_mask, attn_mask_type) and returns: super, super