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","}","$","''","."]
docstr: **Arguments:**
  
  - `num_heads`: Number of parallel attention heads $h$.
  - `query_size`: Number of input channels for query $Q$.
  - `key_size`: Number of input channels for key $K$. Defaults to `query_size`.
  - `value_size`: Number of input channels for value $V$. Defaults to
      `query_size`.
  - `output_size`: Number of output channels. Defaults to `query_size`.
  - `qk_size`: Number of channels to compare query and key over, per head.
      Defaults to `query_size // num_heads`.
  - `vo_size`: Number of channels to compare attention-weighted value and output
      over, per head. Defaults to `query_size // num_heads`.
  - `use_query_bias`: Whether to use a bias term in the query projections.
  - `use_key_bias`: Whether to use a bias term in the key projections.
  - `use_value_bias`: Whether to use a bias term in the value projections.
  - `use_output_bias`: Whether to use a bias term in the output projection.
  - `dropout_p`: Dropout probability on attention weights.
  - `inference`: Whether to actually apply dropout at all. If `True` then dropout
      is not applied. If `False` then dropout is applied. This may be toggled
      with [`equinox.nn.inference_mode`][] or overridden during
      [`equinox.nn.MultiheadAttention.__call__`][].
  - `key`: A `jax.random.PRNGKey` used to provide randomness for parameter
      initialisation. (Keyword only argument.)
docstr_tok: ["*","*","Arguments",":","*","*","-","`","num_heads","`",":","Number","of","parallel","attention","heads","$","h","$",".","-","`","query_size","`",":","Number","of","input","channels","for","query","$","Q","$",".","-","`","key_size","`",":","Number","of","input","channels","for","key","$","K","$",".","Defaults","to","`","query_size","`",".","-","`","value_size","`",":","Number","of","input","channels","for","value","$","V","$",".","Defaults","to","`","query_size","`",".","-","`","output_size","`",":","Number","of","output","channels",".","Defaults","to","`","query_size","`",".","-","`","qk_size","`",":","Number","of","channels","to","compare","query","and","key","over",",","per","head",".","Defaults","to","`","query_size","\/\/","num_heads","`",".","-","`","vo_size","`",":","Number","of","channels","to","compare","attention-weighted","value","and","output","over",",","per","head",".","Defaults","to","`","query_size","\/\/","num_heads","`",".","-","`","use_query_bias","`",":","Whether","to","use","a","bias","term","in","the","query","projections",".","-","`","use_key_bias","`",":","Whether","to","use","a","bias","term","in","the","key","projections",".","-","`","use_value_bias","`",":","Whether","to","use","a","bias","term","in","the","value","projections",".","-","`","use_output_bias","`",":","Whether","to","use","a","bias","term","in","the","output","projection",".","-","`","dropout_p","`",":","Dropout","probability","on","attention","weights",".","-","`","inference","`",":","Whether","to","actually","apply","dropout","at","all",".","If","`","True","`","then","dropout","is","not","applied",".","If","`","False","`","then","dropout","is","applied",".","This","may","be","toggled","with","[","`","equinox.nn.inference_mode","`","]","[","]","or","overridden","during","[","`","equinox.nn.MultiheadAttention.__call__","`","]","[","]",".","-","`","key","`",":","A","`","jax.random.PRNGKey","`","used","to","provide","randomness","for","parameter","initialisation",".","(","Keyword","only","argument",".",")"]
code: def __init__(
      self,
      num_heads: int,
      query_size: int,
      key_size: Optional[int] = None,
      value_size: Optional[int] = None,
      output_size: Optional[int] = None,
      qk_size: Optional[int] = None,
      vo_size: Optional[int] = None,
      use_query_bias: bool = False,
      use_key_bias: bool = False,
      use_value_bias: bool = False,
      use_output_bias: bool = False,
      dropout_p: float = 0.0,
      inference: bool = False,
      *,
      key: PRNGKeyArray,
  ):
      r"""**Arguments:**
  
      - `num_heads`: Number of parallel attention heads $h$.
      - `query_size`: Number of input channels for query $Q$.
      - `key_size`: Number of input channels for key $K$. Defaults to `query_size`.
      - `value_size`: Number of input channels for value $V$. Defaults to
          `query_size`.
      - `output_size`: Number of output channels. Defaults to `query_size`.
      - `qk_size`: Number of channels to compare query and key over, per head.
          Defaults to `query_size // num_heads`.
      - `vo_size`: Number of channels to compare attention-weighted value and output
          over, per head. Defaults to `query_size // num_heads`.
      - `use_query_bias`: Whether to use a bias term in the query projections.
      - `use_key_bias`: Whether to use a bias term in the key projections.
      - `use_value_bias`: Whether to use a bias term in the value projections.
      - `use_output_bias`: Whether to use a bias term in the output projection.
      - `dropout_p`: Dropout probability on attention weights.
      - `inference`: Whether to actually apply dropout at all. If `True` then dropout
          is not applied. If `False` then dropout is applied. This may be toggled
          with [`equinox.nn.inference_mode`][] or overridden during
          [`equinox.nn.MultiheadAttention.__call__`][].
      - `key`: A `jax.random.PRNGKey` used to provide randomness for parameter
          initialisation. (Keyword only argument.)
      """
      qkey, kkey, vkey, okey = jrandom.split(key, 4)
  
      if key_size is None:
          key_size = query_size
      if value_size is None:
          value_size = query_size
      if qk_size is None:
          qk_size = query_size // num_heads
      if vo_size is None:
          vo_size = query_size // num_heads
      if output_size is None:
          output_size = query_size
  
      self.query_proj = Linear(
          query_size,
          num_heads * qk_size,
          use_bias=use_query_bias,
          key=qkey,
      )
      self.key_proj = Linear(
          key_size, num_heads * qk_size, use_bias=use_key_bias, key=kkey
      )
      self.value_proj = Linear(
          value_size,
          num_heads * vo_size,
          use_bias=use_value_bias,
          key=vkey,
      )
      self.output_proj = Linear(
          num_heads * vo_size,
          output_size,
          use_bias=use_output_bias,
          key=okey,
      )
      self.dropout = Dropout(dropout_p, inference=inference)
  
      self.num_heads = num_heads
      self.query_size = query_size
      self.key_size = key_size
      self.value_size = value_size
      self.output_size = output_size
      self.qk_size = qk_size
      self.vo_size = vo_size
      self.use_query_bias = use_query_bias
      self.use_key_bias = use_key_bias
      self.use_value_bias = use_value_bias
      self.use_output_bias = use_output_bias
  
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 1 to create a new MultiheadAttention object from inherited base classes: Module with example: obj = MultiheadAttention(num_heads, query_size, key_size, value_size, output_size, qk_size, vo_size, use_query_bias, use_key_bias, use_value_bias, use_output_bias, dropout_p, inference)