class_docstr: Cross-attention layer class
  
  Cross-attention layer takes input with size [s, b, h] and context with size
  [s, b, h] and returns output of the same size.
class_docstr_tok: ["Cross-attention","layer","class","Cross-attention","layer","takes","input","with","size","[","s",",","b",",","h","]","and","context","with","size","[","s",",","b",",","h","]","and","returns","output","of","the","same","size","."]
docstr: Derives `query` tensor from `hidden_states`, and `key`/`value` tensors
  from `key_value_states`.
docstr_tok: ["Derives","`","query","`","tensor","from","`","hidden_states","`",",","and","`","key","`","\/","`","value","`","tensors","from","`","key_value_states","`","."]
code: def get_query_key_value_tensors(
      self, hidden_states, key_value_states
  ):
      """
      Derives `query` tensor from `hidden_states`, and `key`/`value` tensors
      from `key_value_states`.
      """
      # Attention heads [sk, b, h] --> [sk, b, (np * 2 * hn)]
      mixed_kv, _ = self.linear_kv(key_value_states)
  
      # [sk, b, (np * 2 * hn)] --> [sk, b, np, 2 * hn]
      new_tensor_shape = mixed_kv.size()[:-1] + (
          self.num_attention_heads_per_partition,
          2 * self.hidden_size_per_attention_head,
      )
      mixed_kv = mixed_kv.view(*new_tensor_shape)
  
      # [sk, b, np, 2 * hn] --> 2 [sk, b, np, hn]
      (key, value) = tensor_parallel.split_tensor_along_last_dim(
          mixed_kv, 2
      )
  
      # Attention head [sq, b, h] --> [sq, b, hp]
      query, _ = self.linear_q(hidden_states)
  
      # [sq, b, hp] --> [sq, b, np, hn]
      new_tensor_shape = query.size()[:-1] + (
          self.num_attention_heads_per_partition,
          self.hidden_size_per_attention_head,
      )
      query = query.view(*new_tensor_shape)
  
      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 CrossAttention 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