code: def llama_compute_query_states(
      model: nn.Module, **kwargs
  ) -> torch.Tensor:
      """
      Compute query states for Llama models specifically.
  
      They need to be recomputed as the forward() method of the original LlamaModel in the transformers library does not
      return them. See the related discussion in the PR: https://github.com/huggingface/peft/pull/268
      """
      hidden_states = kwargs.get("hidden_states")
      position_ids = kwargs.get("position_ids")
      past_key_value = kwargs.get("past_key_value")
      bsz, q_len, _ = hidden_states.size()
      query_states = (
          model.q_proj(hidden_states)
          .view(bsz, q_len, model.num_heads, model.head_dim)
          .transpose(1, 2)
      )
      value_states = (
          model.v_proj(hidden_states)
          .view(bsz, q_len, model.num_heads, model.head_dim)
          .transpose(1, 2)
      )
  
      seq_len = q_len
      if past_key_value is not None:
          if isinstance(past_key_value, tuple):
              # for transformers <= 4.35
              seq_len += past_key_value[0].shape[-2]
          else:
              # since transformers 4.36, this is a DynamicCache instance
              seq_len += past_key_value.get_seq_length(model.layer_idx)
      cos, sin = model.rotary_emb(value_states, seq_len=seq_len)
  
      return llama_apply_rotary_pos_emb(
          query_states, cos, sin, position_ids
      )
  
imports: import torch
  import torch.nn
label_desc: Use image node_id 3 for calling a global function with example usage: llama_compute_query_states(model) and returns: llama_apply_rotary_pos_emb