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"]
code: def __init__(self, prefix, config, weights):
      super().__init__()
      self.config = config
      self.embed_dim = config.hidden_size
      self.num_heads = config.num_attention_heads
      self.head_dim = self.embed_dim // self.num_heads
      if self.head_dim * self.num_heads != self.embed_dim:
          raise ValueError(
              f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
              f" {self.num_heads})."
          )
      self.scale = self.head_dim**-0.5
      self.dropout = config.attention_dropout
  
      process_group = weights.process_group
      if self.num_heads % weights.process_group.size() != 0:
          raise ValueError(
              f"`num_heads` must be divisible by `num_shards` (got `num_heads`: {self.num_heads} "
              f"and `num_shards`: {weights.process_group.size()}"
          )
      self.num_heads = self.num_heads // weights.process_group.size()
      self.embed_dim = self.embed_dim // weights.process_group.size()
  
      self.k_proj = TensorParallelColumnLinear.load(
          config, prefix=f"{prefix}.k_proj", weights=weights, bias=True
      )
      self.v_proj = TensorParallelColumnLinear.load(
          config, prefix=f"{prefix}.v_proj", weights=weights, bias=True
      )
      self.q_proj = TensorParallelColumnLinear.load(
          config, prefix=f"{prefix}.q_proj", weights=weights, bias=True
      )
      self.out_proj = TensorParallelRowLinear.load(
          config,
          prefix=f"{prefix}.out_proj",
          weights=weights,
          bias=True,
      )
  
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 1 to create a new IdeficsVisionAttention object from inherited base classes: nn with example: obj = IdeficsVisionAttention(prefix, config, weights)