class_docstr: Multi-head dot-product attention.
  
  Example usage::
  
    >>> import flax.linen as nn
    >>> import jax
  
    >>> layer = nn.MultiHeadDotProductAttention(num_heads=8, qkv_features=16)
    >>> key1, key2, key3, key4, key5, key6 = jax.random.split(jax.random.key(0), 6)
    >>> shape = (4, 3, 2, 5)
    >>> q, k, v = jax.random.uniform(key1, shape), jax.random.uniform(key2, shape), jax.random.uniform(key3, shape)
    >>> variables = layer.init(jax.random.key(0), q)
  
    >>> # different inputs for inputs_q, inputs_k and inputs_v
    >>> out = layer.apply(variables, q, k, v)
    >>> # equivalent to layer.apply(variables, inputs_q=q, inputs_k=k, inputs_v=k)
    >>> out = layer.apply(variables, q, k)
    >>> # equivalent to layer.apply(variables, inputs_q=q, inputs_k=q) and layer.apply(variables, inputs_q=q, inputs_k=q, inputs_v=q)
    >>> out = layer.apply(variables, q)
  
    >>> attention_kwargs = dict(
    ...     num_heads=8,
    ...     qkv_features=16,
    ...     kernel_init=nn.initializers.ones,
    ...     bias_init=nn.initializers.zeros,
    ...     dropout_rate=0.5,
    ...     deterministic=False,
    ...     )
    >>> class Module(nn.Module):
    ...   attention_kwargs: dict
    ...
    ...   @nn.compact
    ...   def __call__(self, x, dropout_rng=None):
    ...     out1 = nn.MultiHeadDotProductAttention(**self.attention_kwargs)(x, dropout_rng=dropout_rng)
    ...     out2 = nn.MultiHeadDotProductAttention(**self.attention_kwargs)(x, dropout_rng=dropout_rng)
    ...     return out1, out2
    >>> module = Module(attention_kwargs)
    >>> variables = module.init({'params': key1, 'dropout': key2}, q)
  
    >>> # out1 and out2 are different.
    >>> out1, out2 = module.apply(variables, q, rngs={'dropout': key3})
    >>> # out3 and out4 are different.
    >>> # out1 and out3 are different. out2 and out4 are different.
    >>> out3, out4 = module.apply(variables, q, rngs={'dropout': key4})
    >>> # out1 and out2 are the same.
    >>> out1, out2 = module.apply(variables, q, dropout_rng=key5)
    >>> # out1 and out2 are the same as out3 and out4.
    >>> # providing a `dropout_rng` arg will take precedence over the `rngs` arg in `.apply`
    >>> out3, out4 = module.apply(variables, q, rngs={'dropout': key6}, dropout_rng=key5)
  
  Attributes:
    num_heads: number of attention heads. Features (i.e. inputs_q.shape[-1])
      should be divisible by the number of heads.
    dtype: the dtype of the computation (default: infer from inputs and params)
    param_dtype: the dtype passed to parameter initializers (default: float32)
    qkv_features: dimension of the key, query, and value.
    out_features: dimension of the last projection
    broadcast_dropout: bool: use a broadcasted dropout along batch dims.
    dropout_rate: dropout rate
    deterministic: if false, the attention weight is masked randomly using
      dropout, whereas if true, the attention weights are deterministic.
    precision: numerical precision of the computation see ``jax.lax.Precision``
      for details.
    kernel_init: initializer for the kernel of the Dense layers.
    bias_init: initializer for the bias of the Dense layers.
    use_bias: bool: whether pointwise QKVO dense transforms use bias.
    attention_fn: dot_product_attention or compatible function. Accepts query,
      key, value, and returns output of shape ``[bs, dim1, dim2, ..., dimN,,
      num_heads, value_channels]``
    decode: whether to prepare and use an autoregressive cache.
    normalize_qk: should QK normalization be applied (arxiv.org/abs/2302.05442).
class_docstr_tok: ["Multi-head","dot-product","attention",".","Example","usage",":",":",">",">",">","import","flax.linen","as","nn",">",">",">","import","jax",">",">",">","layer","=","nn.MultiHeadDotProductAttention","(","num_heads=8",",","qkv_features=16",")",">",">",">","key1",",","key2",",","key3",",","key4",",","key5",",","key6","=","jax.random.split","(","jax.random.key","(","0",")",",","6",")",">",">",">","shape","=","(","4",",","3",",","2",",","5",")",">",">",">","q",",","k",",","v","=","jax.random.uniform","(","key1",",","shape",")",",","jax.random.uniform","(","key2",",","shape",")",",","jax.random.uniform","(","key3",",","shape",")",">",">",">","variables","=","layer.init","(","jax.random.key","(","0",")",",","q",")",">",">",">","#","different","inputs","for","inputs_q",",","inputs_k","and","inputs_v",">",">",">","out","=","layer.apply","(","variables",",","q",",","k",",","v",")",">",">",">","#","equivalent","to","layer.apply","(","variables",",","inputs_q=q",",","inputs_k=k",",","inputs_v=k",")",">",">",">","out","=","layer.apply","(","variables",",","q",",","k",")",">",">",">","#","equivalent","to","layer.apply","(","variables",",","inputs_q=q",",","inputs_k=q",")","and","layer.apply","(","variables",",","inputs_q=q",",","inputs_k=q",",","inputs_v=q",")",">",">",">","out","=","layer.apply","(","variables",",","q",")",">",">",">","attention_kwargs","=","dict","(","...","num_heads=8",",","...","qkv_features=16",",","...","kernel_init=nn.initializers.ones",",","...","bias_init=nn.initializers.zeros",",","...","dropout_rate=0.5",",","...","deterministic=False",",","...",")",">",">",">","class","Module","(","nn.Module",")",":","...","attention_kwargs",":","dict","...","...","@","nn.compact","...","def","__call__","(","self",",","x",",","dropout_rng=None",")",":","...","out1","=","nn.MultiHeadDotProductAttention","(","*","*","self.attention_kwargs",")","(","x",",","dropout_rng=dropout_rng",")","...","out2","=","nn.MultiHeadDotProductAttention","(","*","*","self.attention_kwargs",")","(","x",",","dropout_rng=dropout_rng",")","...","return","out1",",","out2",">",">",">","module","=","Module","(","attention_kwargs",")",">",">",">","variables","=","module.init","(","{","'params","'",":","key1",",","'dropout","'",":","key2","}",",","q",")",">",">",">","#","out1","and","out2","are","different",".",">",">",">","out1",",","out2","=","module.apply","(","variables",",","q",",","rngs=","{","'dropout","'",":","key3","}",")",">",">",">","#","out3","and","out4","are","different",".",">",">",">","#","out1","and","out3","are","different",".","out2","and","out4","are","different",".",">",">",">","out3",",","out4","=","module.apply","(","variables",",","q",",","rngs=","{","'dropout","'",":","key4","}",")",">",">",">","#","out1","and","out2","are","the","same",".",">",">",">","out1",",","out2","=","module.apply","(","variables",",","q",",","dropout_rng=key5",")",">",">",">","#","out1","and","out2","are","the","same","as","out3","and","out4",".",">",">",">","#","providing","a","`","dropout_rng","`","arg","will","take","precedence","over","the","`","rngs","`","arg","in","`",".apply","`",">",">",">","out3",",","out4","=","module.apply","(","variables",",","q",",","rngs=","{","'dropout","'",":","key6","}",",","dropout_rng=key5",")","Attributes",":","num_heads",":","number","of","attention","heads",".","Features","(","i.e",".","inputs_q.shape","[","-1","]",")","should","be","divisible","by","the","number","of","heads",".","dtype",":","the","dtype","of","the","computation","(","default",":","infer","from","inputs","and","params",")","param_dtype",":","the","dtype","passed","to","parameter","initializers","(","default",":","float32",")","qkv_features",":","dimension","of","the","key",",","query",",","and","value",".","out_features",":","dimension","of","the","last","projection","broadcast_dropout",":","bool",":","use","a","broadcasted","dropout","along","batch","dims",".","dropout_rate",":","dropout","rate","deterministic",":","if","false",",","the","attention","weight","is","masked","randomly","using","dropout",",","whereas","if","true",",","the","attention","weights","are","deterministic",".","precision",":","numerical","precision","of","the","computation","see","``","jax.lax.Precision","``","for","details",".","kernel_init",":","initializer","for","the","kernel","of","the","Dense","layers",".","bias_init",":","initializer","for","the","bias","of","the","Dense","layers",".","use_bias",":","bool",":","whether","pointwise","QKVO","dense","transforms","use","bias",".","attention_fn",":","dot_product_attention","or","compatible","function",".","Accepts","query",",","key",",","value",",","and","returns","output","of","shape","``","[","bs",",","dim1",",","dim2",",","...",",","dimN",",",",","num_heads",",","value_channels","]","``","decode",":","whether","to","prepare","and","use","an","autoregressive","cache",".","normalize_qk",":","should","QK","normalization","be","applied","(","arxiv.org\/abs\/2302.05442",")","."]
code: def __call__(
      self,
      inputs_q: Array,
      *,
      inputs_kv: Array = None,
      mask: Optional[Array] = None,
      deterministic: Optional[bool] = None,
      dropout_rng: Optional[PRNGKey] = None,
      sow_weights: bool = False,
  ):
      ...
  
imports: import functools
  import warnings
  from typing import Any, Callable, Optional, Tuple, Union, overload
  import jax
  import jax.numpy
  from jax import lax, random
  from flax.linen import initializers
  from flax.linen.dtypes import promote_dtype
  from flax.linen.linear import DenseGeneral, DotGeneralT, PrecisionLike, default_kernel_init
  from flax.linen.module import Module, compact, merge_param
  from flax.linen.normalization import LayerNorm
label_desc: Use image node_id 2 for calling the MultiHeadDotProductAttention obj's underlying member method code with example usage: obj.__call__(inputs_q) without return types