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:**
  
  - `query`: Query embedding. Should be a JAX array of shape
      `(query_seq_length, query_size)`.
  - `key_`: Key embedding. Should be a JAX array of shape
      `(kv_seq_length, key_size)`.
  - `value`: Value embedding. Should be a JAX array of shape
      `(kv_seq_length, value_size)`.
  - `mask`: Optional mask preventing attention to certain positions. Should either
      be a JAX array of shape `(query_seq_length, kv_seq_length)`, or (for custom
      per-head masking) `(num_heads, query_seq_length, kv_seq_length)`. A value of
      `False` at a position indicates that position should be ignored.
  - `key`: A `jax.random.PRNGKey` used for dropout. Unused if `dropout = 0`.
      (Keyword only argument.)
  - `inference`: As [`equinox.nn.Dropout.__call__`][]. (Keyword only
      argument.)
  - `deterministic`: (Deprecated in favour of `inference`.)
  
  **Returns:**
  
  A JAX array of shape `(query_seq_length, output_size)`.
docstr_tok: ["*","*","Arguments",":","*","*","-","`","query","`",":","Query","embedding",".","Should","be","a","JAX","array","of","shape","`","(","query_seq_length",",","query_size",")","`",".","-","`","key_","`",":","Key","embedding",".","Should","be","a","JAX","array","of","shape","`","(","kv_seq_length",",","key_size",")","`",".","-","`","value","`",":","Value","embedding",".","Should","be","a","JAX","array","of","shape","`","(","kv_seq_length",",","value_size",")","`",".","-","`","mask","`",":","Optional","mask","preventing","attention","to","certain","positions",".","Should","either","be","a","JAX","array","of","shape","`","(","query_seq_length",",","kv_seq_length",")","`",",","or","(","for","custom","per-head","masking",")","`","(","num_heads",",","query_seq_length",",","kv_seq_length",")","`",".","A","value","of","`","False","`","at","a","position","indicates","that","position","should","be","ignored",".","-","`","key","`",":","A","`","jax.random.PRNGKey","`","used","for","dropout",".","Unused","if","`","dropout","=","0","`",".","(","Keyword","only","argument",".",")","-","`","inference","`",":","As","[","`","equinox.nn.Dropout.__call__","`","]","[","]",".","(","Keyword","only","argument",".",")","-","`","deterministic","`",":","(","Deprecated","in","favour","of","`","inference","`",".",")","*","*","Returns",":","*","*","A","JAX","array","of","shape","`","(","query_seq_length",",","output_size",")","`","."]
code: def __call__(
      self,
      query: Float[Array, "q_seq q_size"],
      key_: Float[Array, "kv_seq k_size"],
      value: Float[Array, "kv_seq v_size"],
      mask: Union[
          None,
          Bool[Array, "q_seq kv_seq"],
          Bool[Array, "num_heads q_seq kv_seq"],
      ] = None,
      *,
      key: Optional[PRNGKeyArray] = None,
      inference: Optional[bool] = None,
      deterministic: Optional[bool] = None,
  ) -> Float[Array, "q_seq o_size"]:
      """**Arguments:**
  
      - `query`: Query embedding. Should be a JAX array of shape
          `(query_seq_length, query_size)`.
      - `key_`: Key embedding. Should be a JAX array of shape
          `(kv_seq_length, key_size)`.
      - `value`: Value embedding. Should be a JAX array of shape
          `(kv_seq_length, value_size)`.
      - `mask`: Optional mask preventing attention to certain positions. Should either
          be a JAX array of shape `(query_seq_length, kv_seq_length)`, or (for custom
          per-head masking) `(num_heads, query_seq_length, kv_seq_length)`. A value of
          `False` at a position indicates that position should be ignored.
      - `key`: A `jax.random.PRNGKey` used for dropout. Unused if `dropout = 0`.
          (Keyword only argument.)
      - `inference`: As [`equinox.nn.Dropout.__call__`][]. (Keyword only
          argument.)
      - `deterministic`: (Deprecated in favour of `inference`.)
  
      **Returns:**
  
      A JAX array of shape `(query_seq_length, output_size)`.
      """
  
      if deterministic is not None:
          inference = deterministic
          warnings.warn(
              "MultiheadAttention()(deterministic=...) is deprecated "
              "in favour of MultiheadAttention()(inference=...)"
          )
  
      query_seq_length, _ = query.shape
      kv_seq_length, _ = key_.shape
      kv_seq_length2, _ = value.shape
      if kv_seq_length != kv_seq_length2:
          # query length can be different
          raise ValueError(
              "key and value must both be sequences of equal length."
          )
  
      query_heads = self._project(self.query_proj, query)
      key_heads = self._project(self.key_proj, key_)
      value_heads = self._project(self.value_proj, value)
  
      attn_fn = partial(
          dot_product_attention,
          dropout=self.dropout,
          inference=inference,
      )
      keys = (
          None
          if key is None
          else jax.random.split(key, query_heads.shape[1])
      )
      if mask is not None and mask.ndim == 3:
          # Batch `mask` and `keys` down their 0-th dimension.
          attn = jax.vmap(attn_fn, in_axes=1, out_axes=1)(
              query_heads, key_heads, value_heads, mask=mask, key=keys
          )
      else:
          # Batch `keys` down its 0-th dimension.
          attn = jax.vmap(
              ft.partial(attn_fn, mask=mask), in_axes=1, out_axes=1
          )(query_heads, key_heads, value_heads, key=keys)
      attn = attn.reshape(query_seq_length, -1)
  
      return jax.vmap(self.output_proj)(attn)
  
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 2 for calling the MultiheadAttention obj's underlying member method code with example usage: obj.__call__(query, key_, value, mask) and returns: jax