class_docstr: Self-attention layer class
  
  Self-attention layer takes input with size [s, b, h]
  and returns output of the same size.
class_docstr_tok: ["Self-attention","layer","class","Self-attention","layer","takes","input","with","size","[","s",",","b",",","h","]","and","returns","output","of","the","same","size","."]
docstr: Derives `query`, `key` and `value` tensors from `hidden_states`.
docstr_tok: ["Derives","`","query","`",",","`","key","`","and","`","value","`","tensors","from","`","hidden_states","`","."]
code: def get_query_key_value_tensors(
      self, hidden_states, key_value_states=None
  ):
      """
      Derives `query`, `key` and `value` tensors from `hidden_states`.
      """
      # Attention heads [sq, b, h] --> [sq, b, ng * (np/ng + 2) * hn)]
      mixed_qkv, _ = self.linear_qkv(hidden_states)
  
      # [sq, b, hp] --> [sq, b, ng, (np/ng + 2) * hn]
      new_tensor_shape = mixed_qkv.size()[:-1] + (
          self.num_query_groups_per_partition,
          (
              (
                  self.num_attention_heads_per_partition
                  // self.num_query_groups_per_partition
                  + 2
              )
              * self.hidden_size_per_attention_head
          ),
      )
      mixed_qkv = mixed_qkv.view(*new_tensor_shape)
  
      split_arg_list = [
          (
              self.num_attention_heads_per_partition
              // self.num_query_groups_per_partition
              * self.hidden_size_per_attention_head
          ),
          self.hidden_size_per_attention_head,
          self.hidden_size_per_attention_head,
      ]
  
      if SplitAlongDim is not None:
          # [sq, b, ng, (np/ng + 2) * hn] --> [sq, b, ng, np/ng * hn], [sq, b, ng, hn], [sq, b, ng, hn]
          (query, key, value) = SplitAlongDim(
              mixed_qkv,
              3,
              split_arg_list,
          )
      else:
          # [sq, b, ng, (np/ng + 2) * hn] --> [sq, b, ng, np/ng * hn], [sq, b, ng, hn], [sq, b, ng, hn]
          (query, key, value) = torch.split(
              mixed_qkv,
              split_arg_list,
              dim=3,
          )
  
      # [sq, b, ng, np/ng * hn] -> [sq, b, np, hn]
      query = query.reshape(
          query.size(0),
          query.size(1),
          -1,
          self.hidden_size_per_attention_head,
      )
  
      return query, key, value
  
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 SelfAttention obj's underlying member method code with example usage: obj.get_query_key_value_tensors(hidden_states, key_value_states) and returns: query,  key,  value