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)]
      linear_qkv_output, _ = self.linear_qkv(hidden_states)
      layernorm_output = None
  
      # In megatron/core/models/gpt/gpt_layer_specs.py TELayerNormColumnParallelLinear is used for linear_qkv.
      # TELayerNormColumnParallelLinear fused LN and linear, both will be returned.
      # In nemo/collections/nlp/models/language_modeling/megatron/falcon/falcon_spec.py TEColumnParallelLinear is used for linear_qkv,
      # which only returns linear.
      if isinstance(self.linear_qkv, TELayerNormColumnParallelLinear):
          mixed_qkv, layernorm_output = linear_qkv_output
      elif isinstance(
          self.linear_qkv, TEColumnParallelLinear
      ):  # only mixed_qkv
          mixed_qkv = linear_qkv_output
      else:
          raise ValueError(
              f"Unrecognized module type '{type(self.linear_qkv)}' when getting query, key, value tensors for mcore mixins. "
          )
  
      # LoRA logic
      if self.is_adapter_available():
          lora_kqv_adapter = self.get_adapter_module(
              AdapterName.LORA_KQV_ADAPTER
          )
          if lora_kqv_adapter:
              if isinstance(
                  self.linear_qkv, TELayerNormColumnParallelLinear
              ):
                  lora_mixed_qkv = lora_kqv_adapter(layernorm_output)
              elif isinstance(self.linear_qkv, TEColumnParallelLinear):
                  lora_mixed_qkv = lora_kqv_adapter(hidden_states)
              else:
                  raise ValueError(
                      f"Unrecognized module type '{type(self.linear_qkv)}' when applying lora."
                  )
              mixed_qkv = mixed_qkv + lora_mixed_qkv
  
      # [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)
  
      # [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,
          [
              (
                  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,
          ],
          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,
      )
  
      if self.is_adapter_available():
          key_infused_adapter = self.get_adapter_module(
              AdapterName.KEY_INFUSED
          )
          value_infused_adapter = self.get_adapter_module(
              AdapterName.VALUE_INFUSED
          )
          if key_infused_adapter:
              assert (
                  value_infused_adapter is not None
              ), "Expected value_infused_adapter not found!"
              kls = key.shape
              key = (
                  key_infused_adapter(key.reshape(kls[0], kls[1], -1))
                  .reshape(kls)
                  .to(query.dtype)
              )
          if value_infused_adapter:
              assert (
                  key_infused_adapter is not None
              ), "Expected key_infused_adapter not found!"
              vls = value.shape
              value = (
                  value_infused_adapter(
                      value.reshape(vls[0], vls[1], -1)
                  )
                  .reshape(vls)
                  .to(query.dtype)
              )
  
      return query, key, value
  
imports: import torch
  import torch.nn.functional
  from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
  from megatron.core.fusions.fused_bias_gelu import bias_gelu_impl
  from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding
  from megatron.core.transformer.attention import SelfAttention
  from megatron.core.transformer.custom_layers.transformer_engine import TEColumnParallelLinear, TELayerNormColumnParallelLinear
  from megatron.core.transformer.mlp import MLP
  from megatron.core.transformer.transformer_layer import TransformerLayer
  from megatron.core.utils import make_viewless_tensor
  from nemo.collections.nlp.modules.common.megatron.adapters.parallel_adapters import AdapterName, InfusedAdapterConfig, LoraKQVAdapterConfig, MLPInfusedAdapterConfig, ParallelLinearAdapterConfig, PromptEncoderAdapterConfig
  from nemo.core import adapter_mixins
label_desc: Use image node_id 2 for calling the MCoreSelfAttentionMixin 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