docstr: Perceiver Cross-Attention Module --> let long-form inputs be `context`, resampled embeddings be `latents`
docstr_tok: ["Perceiver","Cross-Attention","Module","--",">","let","long-form","inputs","be","`","context","`",",","resampled","embeddings","be","`","latents","`"]
code: def __init__(
      self,
      prefix,
      config,
      embed_dim: int,
      n_heads: int,
      head_dim: int,
      qk_layer_norms: bool,
      weights,
  ) -> None:
      """Perceiver Cross-Attention Module --> let long-form inputs be `context`, resampled embeddings be `latents`"""
      super().__init__()
      self.embed_dim, self.n_heads, self.head_dim = (
          embed_dim,
          n_heads,
          head_dim,
      )
      self.qk_layer_norms = qk_layer_norms
      # Normalization & Scaling
      self.context_layer_norm = nn.LayerNorm.load(
          prefix=f"{prefix}.context_layer_norm",
          weights=weights,
          eps=EPS,
      )
      self.latents_layer_norm = nn.LayerNorm.load(
          prefix=f"{prefix}.latents_layer_norm",
          weights=weights,
          eps=EPS,
      )
      if self.qk_layer_norms:
          self.q_layer_norm = nn.LayerNorm.load(
              prefix=f"{prefix}.q_layer_norm", weights=weights, eps=EPS
          )
          self.k_layer_norm = nn.LayerNorm.load(
              prefix=f"{prefix}.k_layer_norm", weights=weights, eps=EPS
          )
  
      self.qk_scale = self.head_dim**-0.5
  
      process_group = weights.process_group
      if n_heads % weights.process_group.size() != 0:
          raise ValueError(
              f"`num_heads` must be divisible by `num_shards` (got `num_heads`: {n_heads} "
              f"and `num_shards`: {weights.process_group.size()}"
          )
      self.n_heads //= weights.process_group.size()
  
      # Q, K, V Projection (no bias -- detail from Perceiver/Flamingo Papers).
      self.q_proj = TensorParallelColumnLinear.load(
          config=config,
          prefix=f"{prefix}.q_proj",
          weights=weights,
          bias=False,
      )
      self.k_proj = TensorParallelColumnLinear.load(
          config=config,
          prefix=f"{prefix}.k_proj",
          weights=weights,
          bias=False,
      )
      self.v_proj = TensorParallelColumnLinear.load(
          config=config,
          prefix=f"{prefix}.v_proj",
          weights=weights,
          bias=False,
      )
  
      self.output_proj = TensorParallelRowLinear.load(
          config=config,
          prefix=f"{prefix}.output_proj",
          weights=weights,
          bias=False,
      )
  
file_path: text-generation-inference/server/text_generation_server/models/custom_modeling/idefics_perceiver.py
imports: from typing import Optional, Tuple
  import torch
  import torch.nn
  from text_generation_server.utils.layers import TensorParallelColumnLinear, TensorParallelRowLinear
label_desc: Use image node_id 1 to create a new IdeficsPerceiverAttention object from inherited base classes: nn with example: obj = IdeficsPerceiverAttention(prefix, config, embed_dim, n_heads, head_dim, qk_layer_norms, weights)