class_docstr: Multi-head attention.
  
  Example usage::
  
    >>> import flax.linen as nn
    >>> import jax
  
    >>> layer = nn.MultiHeadAttention(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.MultiHeadAttention(**self.attention_kwargs)(x, dropout_rng=dropout_rng)
    ...     out2 = nn.MultiHeadAttention(**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","attention",".","Example","usage",":",":",">",">",">","import","flax.linen","as","nn",">",">",">","import","jax",">",">",">","layer","=","nn.MultiHeadAttention","(","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.MultiHeadAttention","(","*","*","self.attention_kwargs",")","(","x",",","dropout_rng=dropout_rng",")","...","out2","=","nn.MultiHeadAttention","(","*","*","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 = None,
      mask: Array | None = None,
      deterministic: bool | None = None,
      dropout_rng: Array | None = None,
      rngs: rnglib.Rngs | None = None,
      sow_weights: bool = False,
      decode: bool | None = None,
  ):
      ...
  
imports: from __future__ import annotations
  import functools
  from typing import Any, Callable, Optional, Tuple, overload
  import jax
  import jax.numpy
  from jax import lax, random
  from flax.experimental import nnx
  from flax.experimental.nnx.nnx import rnglib
  from flax.experimental.nnx.nnx import flaglib
  from flax.experimental.nnx.nnx.module import Module, first_from
  from flax.experimental.nnx.nnx.nn import initializers
  from flax.experimental.nnx.nnx.nn.dtypes import promote_dtype
  from flax.experimental.nnx.nnx.nn.linear import DotGeneralT, LinearGeneral, PrecisionLike, default_kernel_init
  from flax.experimental.nnx.nnx.nn.normalization import LayerNorm
label_desc: Use image node_id 3 for calling the MultiHeadAttention obj's underlying member method code with example usage: obj.__call__(inputs_q) without return types