class_docstr: Multi-headed attention from 'Attention Is All You Need' paper
code: def forward(
      self,
      hidden_states: torch.Tensor,
      attention_mask: Optional[torch.Tensor] = None,
      position_ids: Optional[torch.LongTensor] = None,
      **kwargs,
  ) -> torch.Tensor:
      bsz, q_len, _ = hidden_states.size()
  
      query_states = self.q_proj(hidden_states)
      key_states = self.k_proj(hidden_states)
      value_states = self.v_proj(hidden_states)
  
      query_states = query_states.view(
          bsz, q_len, self.num_heads, self.head_dim
      ).transpose(1, 2)
      key_states = key_states.view(
          bsz, q_len, self.num_heads, self.head_dim
      ).transpose(1, 2)
      value_states = value_states.view(
          bsz, q_len, self.num_heads, self.head_dim
      ).transpose(1, 2)
  
      kv_seq_len = key_states.shape[-2]
      cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
      query_states, key_states = apply_rotary_pos_emb(
          query_states, key_states, cos, sin, position_ids
      )
  
      attn_weights = torch.matmul(
          query_states, key_states.transpose(2, 3)
      ) / math.sqrt(self.head_dim)
  
      if attention_mask is not None:
          if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
              raise ValueError(
                  f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
              )
          attn_weights = attn_weights + attention_mask
  
      # upcast attention to fp32
      attn_weights = nn.functional.softmax(
          attn_weights, dim=-1, dtype=torch.float32
      ).to(query_states.dtype)
      attn_output = torch.matmul(attn_weights, value_states)
  
      attn_output = attn_output.transpose(1, 2).contiguous()
  
      attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
  
      return self.o_proj(attn_output)
  
imports: import math
  from typing import Optional
  import pytest
  import torch
  import torch.utils.checkpoint
  from helpers import assert_similar, random_tensor
  from torch import nn
  from quanto import Calibration, quantize
label_desc: Use image node_id 3 for calling the Attention obj's underlying member method code with example usage: obj.forward(hidden_states, attention_mask, position_ids) and returns: self