class_docstr: Multi-headed attention from 'Attention Is All You Need' paper
class_docstr_tok: ["Multi-headed","attention","from","'Attention","Is","All","You","Need","'","paper"]
docstr: Input shape: Batch x Time x Channel
docstr_tok: ["Input","shape",":","Batch","x","Time","x","Channel"]
code: def forward(
      self,
      hidden_states: torch.Tensor,
      attention_mask: Optional[torch.Tensor] = None,
      causal_attention_mask: Optional[torch.Tensor] = None,
      output_attentions: Optional[bool] = False,
  ) -> Tuple[
      torch.Tensor,
      Optional[torch.Tensor],
      Optional[Tuple[torch.Tensor]],
  ]:
      """Input shape: Batch x Time x Channel"""
  
      bsz, tgt_len, _ = hidden_states.size()
  
      # get query proj
      query_states = self.q_proj(hidden_states) * self.scale
      key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
      value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
  
      proj_shape = (bsz * self.num_heads, -1, self.head_dim)
      query_states = self._shape(query_states, tgt_len, bsz).view(
          *proj_shape
      )
      key_states = key_states.view(*proj_shape)
      value_states = value_states.view(*proj_shape)
  
      src_len = key_states.size(1)
      attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))
  
      if attn_weights.size() != (
          bsz * self.num_heads,
          tgt_len,
          src_len,
      ):
          raise ValueError(
              f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is"
              f" {attn_weights.size()}"
          )
  
      # apply the causal_attention_mask first
      if causal_attention_mask is not None:
          if causal_attention_mask.size() != (bsz, 1, tgt_len, src_len):
              raise ValueError(
                  f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is"
                  f" {causal_attention_mask.size()}"
              )
          attn_weights = (
              attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
              + causal_attention_mask
          )
          attn_weights = attn_weights.view(
              bsz * self.num_heads, tgt_len, src_len
          )
  
      if attention_mask is not None:
          if attention_mask.size() != (bsz, 1, tgt_len, src_len):
              raise ValueError(
                  f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}"
              )
          attn_weights = (
              attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
              + attention_mask
          )
          attn_weights = attn_weights.view(
              bsz * self.num_heads, tgt_len, src_len
          )
  
      attn_weights = nn.functional.softmax(attn_weights, dim=-1)
  
      if output_attentions:
          # this operation is a bit akward, but it's required to
          # make sure that attn_weights keeps its gradient.
          # In order to do so, attn_weights have to reshaped
          # twice and have to be reused in the following
          attn_weights_reshaped = attn_weights.view(
              bsz, self.num_heads, tgt_len, src_len
          )
          attn_weights = attn_weights_reshaped.view(
              bsz * self.num_heads, tgt_len, src_len
          )
      else:
          attn_weights_reshaped = None
  
      attn_probs = nn.functional.dropout(
          attn_weights, p=self.dropout, training=self.training
      )
  
      attn_output = torch.bmm(attn_probs, value_states)
  
      if attn_output.size() != (
          bsz * self.num_heads,
          tgt_len,
          self.head_dim,
      ):
          raise ValueError(
              f"`attn_output` should be of size {(bsz, self.num_heads, tgt_len, self.head_dim)}, but is"
              f" {attn_output.size()}"
          )
  
      attn_output = attn_output.view(
          bsz, self.num_heads, tgt_len, self.head_dim
      )
      attn_output = attn_output.transpose(1, 2)
      attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim)
  
      attn_output = self.out_proj(attn_output)
  
      return attn_output, attn_weights_reshaped
  
imports: from dataclasses import dataclass
  from typing import Optional, Tuple, Union
  import torch
  import torch.utils.checkpoint
  from torch import nn
  from transformers.activations import ACT2FN
  from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling
  from transformers.utils import ModelOutput, logging
  from text_generation_server.utils.layers import TensorParallelColumnLinear, TensorParallelRowLinear, TensorParallelEmbedding
label_desc: Use image node_id 3 for calling the IdeficsVisionAttention obj's underlying member method code with example usage: obj.forward(hidden_states, attention_mask, causal_attention_mask, output_attentions) and returns: attn_output,  attn_weights_reshaped