class_docstr: Computes
  
  $$\text{MultiheadAttention}(Q, K, V)
    = \sum_i \text{Attention}\left(QW^Q_i, KW^K_i, VW^V_i\right)W^O_i$$
  
  where:
  
  - The inputs are
    $Q \in \mathbb{R}^{d_\text{seq} \times d_\text{query}}$,
    $K \in \mathbb{R}^{d_\text{seq} \times d_\text{key}}$,
    $V \in \mathbb{R}^{d_\text{seq} \times d_\text{value}}$.
    These are referred to as query, key, and value respectively. Meanwhile
    $d_\text{seq}$ is the sequence length, and $d_\text{query}$, $d_\text{key}$,
    $d_\text{value}$ are numbers of channels.
  
  - The trainable weights are
  $W^Q_i \in \mathbb{R}^{d_\text{query} \times d_\text{qk}}$,
  $W^K_i \in \mathbb{R}^{d_\text{key} \times d_\text{qk}}$,
  $W^V_i \in \mathbb{R}^{d_\text{value} \times d_\text{vo}}$,
  $W^O_i \in \mathbb{R}^{d_\text{vo} \times d_\text{output}}$,
  with $i \in \{1, \ldots, h\}$, where $h$ is the number of heads, and $d_\text{qk}$,
  $d_\text{vo}$, $d_\text{output}$ are hyperparameters.
  
  - $\text{Attention}$ is defined as
    $\text{Attention}(\widetilde{Q}, \widetilde{K}, \widetilde{V})
     = \text{softmax}(\frac{\widetilde{Q}\widetilde{K}^\intercal}
                           {\sqrt{d_\text{qk}}})\widetilde{V}$.
  
  ??? cite
  
      [Attention is All You Need](https://arxiv.org/abs/1706.03762)
  
      ```bibtex
      @inproceedings{vaswani2017attention,
          author={Vaswani, Ashish and Shazeer, Noam and Parmar, Niki and
                  Uszkoreit, Jakob and Jones, Llion and Gomez, Aidan N and
                  Kaiser, {\L}ukasz and Polosukhin, Illia},
          booktitle={Advances in Neural Information Processing Systems},
          publisher={Curran Associates, Inc.},
          title={Attention is All You Need},
          volume={30},
          year={2017}
      }
      ```
  
  !!! faq "FAQ"
  
      Different software libraries often implement multihead attention in slightly
      different ways. Some of them will or won't add on biases by default. Most of
      them will fix the values of $d_\text{qk}, d_\text{vo}, d_\text{output}$ in
      terms of $d_\text{query}$ or $d_\text{key}$ or $d_\text{value}$. Equinox
      chooses to expose all of these as options.
  
      Relative to the original
      [Attention is All You Need](https://arxiv.org/abs/1706.03762) paper: our
      $d_\text{qk}$ is their "$d_k$". Our $d_\text{vo}$ is their "$d_\text{v}$". They
      fix $d_\text{query} = d_\text{key} = d_\text{value} = d_\text{output}$ and
      refer to it as "$d_\text{model}$".
class_docstr_tok: ["Computes","$","$","\\text","{","MultiheadAttention","}","(","Q",",","K",",","V",")","=","\\sum_i","\\text","{","Attention","}","\\left","(","QW^Q_i",",","KW^K_i",",","VW^V_i\\right",")","W^O_i","$","$","where",":","-","The","inputs","are","$","Q","\\in","\\mathbb","{","R","}","^","{","d_\\text","{","seq","}","\\times","d_\\text","{","query","}","}","$",",","$","K","\\in","\\mathbb","{","R","}","^","{","d_\\text","{","seq","}","\\times","d_\\text","{","key","}","}","$",",","$","V","\\in","\\mathbb","{","R","}","^","{","d_\\text","{","seq","}","\\times","d_\\text","{","value","}","}","$",".","These","are","referred","to","as","query",",","key",",","and","value","respectively",".","Meanwhile","$","d_\\text","{","seq","}","$","is","the","sequence","length",",","and","$","d_\\text","{","query","}","$",",","$","d_\\text","{","key","}","$",",","$","d_\\text","{","value","}","$","are","numbers","of","channels",".","-","The","trainable","weights","are","$","W^Q_i","\\in","\\mathbb","{","R","}","^","{","d_\\text","{","query","}","\\times","d_\\text","{","qk","}","}","$",",","$","W^K_i","\\in","\\mathbb","{","R","}","^","{","d_\\text","{","key","}","\\times","d_\\text","{","qk","}","}","$",",","$","W^V_i","\\in","\\mathbb","{","R","}","^","{","d_\\text","{","value","}","\\times","d_\\text","{","vo","}","}","$",",","$","W^O_i","\\in","\\mathbb","{","R","}","^","{","d_\\text","{","vo","}","\\times","d_\\text","{","output","}","}","$",",","with","$","i","\\in","\\","{","1",",","\\ldots",",","h\\","}","$",",","where","$","h","$","is","the","number","of","heads",",","and","$","d_\\text","{","qk","}","$",",","$","d_\\text","{","vo","}","$",",","$","d_\\text","{","output","}","$","are","hyperparameters",".","-","$","\\text","{","Attention","}","$","is","defined","as","$","\\text","{","Attention","}","(","\\widetilde","{","Q","}",",","\\widetilde","{","K","}",",","\\widetilde","{","V","}",")","=","\\text","{","softmax","}","(","\\frac","{","\\widetilde","{","Q","}","\\widetilde","{","K","}","^\\intercal","}","{","\\sqrt","{","d_\\text","{","qk","}","}","}",")","\\widetilde","{","V","}","$",".","?","?","?","cite","[","Attention","is","All","You","Need","]","(","https",":","\/\/arxiv.org\/abs\/1706.03762",")","``","`","bibtex","@","inproceedings","{","vaswani2017attention",",","author=","{","Vaswani",",","Ashish","and","Shazeer",",","Noam","and","Parmar",",","Niki","and","Uszkoreit",",","Jakob","and","Jones",",","Llion","and","Gomez",",","Aidan","N","and","Kaiser",",","{","\\L","}","ukasz","and","Polosukhin",",","Illia","}",",","booktitle=","{","Advances","in","Neural","Information","Processing","Systems","}",",","publisher=","{","Curran","Associates",",","Inc.","}",",","title=","{","Attention","is","All","You","Need","}",",","volume=","{","30","}",",","year=","{","2017","}","}","``","`","!","!","!","faq","``","FAQ","''","Different","software","libraries","often","implement","multihead","attention","in","slightly","different","ways",".","Some","of","them","will","or","wo","n't","add","on","biases","by","default",".","Most","of","them","will","fix","the","values","of","$","d_\\text","{","qk","}",",","d_\\text","{","vo","}",",","d_\\text","{","output","}","$","in","terms","of","$","d_\\text","{","query","}","$","or","$","d_\\text","{","key","}","$","or","$","d_\\text","{","value","}","$",".","Equinox","chooses","to","expose","all","of","these","as","options",".","Relative","to","the","original","[","Attention","is","All","You","Need","]","(","https",":","\/\/arxiv.org\/abs\/1706.03762",")","paper",":","our","$","d_\\text","{","qk","}","$","is","their","``","$","d_k","$","''",".","Our","$","d_\\text","{","vo","}","$","is","their","``","$","d_\\text","{","v","}","$","''",".","They","fix","$","d_\\text","{","query","}","=","d_\\text","{","key","}","=","d_\\text","{","value","}","=","d_\\text","{","output","}","$","and","refer","to","it","as","``","$","d_\\text","{","model","}","$","''","."]
code: def _project(self, proj, x):
      seq_length, _ = x.shape
      projection = jax.vmap(proj)(x)
      return projection.reshape(seq_length, self.num_heads, -1)
  
file_path: equinox/equinox/nn/_attention.py
imports: import functools
  import math
  import warnings
  from functools import partial
  from typing import cast, Optional, Union
  import jax
  import jax.numpy
  import jax.random
  from jaxtyping import Array, Bool, Float, PRNGKeyArray
  from .._module import field, Module
  from ._dropout import Dropout
  from ._linear import Linear
label_desc: Use image node_id 3 for calling the MultiheadAttention obj's underlying member method code with example usage: obj._project(proj, x) and returns: projection