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 __init__(
      self,
      num_heads: int,
      in_features: int,
      qkv_features: int | None = None,
      out_features: int | None = None,
      *,
      dtype: Dtype | None = None,
      param_dtype: Dtype = jnp.float32,
      broadcast_dropout: bool = True,
      dropout_rate: float = 0.0,
      deterministic: bool | None = None,
      precision: PrecisionLike = None,
      kernel_init: initializers.Initializer = default_kernel_init,
      bias_init: initializers.Initializer = initializers.zeros_init(),
      use_bias: bool = True,
      attention_fn: Callable[..., Array] = dot_product_attention,
      decode: bool | None = None,
      normalize_qk: bool = False,
      # Deprecated, will be removed.
      qkv_dot_general: DotGeneralT | None = None,
      out_dot_general: DotGeneralT | None = None,
      qkv_dot_general_cls: Any = None,
      out_dot_general_cls: Any = None,
      rngs: rnglib.Rngs,
  ):
      self.num_heads = num_heads
      self.in_features = in_features
      self.qkv_features = (
          qkv_features if qkv_features is not None else in_features
      )
      self.out_features = (
          out_features if out_features is not None else in_features
      )
      self.dtype = dtype
      self.param_dtype = param_dtype
      self.broadcast_dropout = broadcast_dropout
      self.dropout_rate = dropout_rate
      self.deterministic = deterministic
      self.precision = precision
      self.kernel_init = kernel_init
      self.bias_init = bias_init
      self.use_bias = use_bias
      self.attention_fn = attention_fn
      self.decode = decode
      self.normalize_qk = normalize_qk
      self.qkv_dot_general = qkv_dot_general
      self.out_dot_general = out_dot_general
      self.qkv_dot_general_cls = qkv_dot_general_cls
      self.out_dot_general_cls = out_dot_general_cls
  
      if self.qkv_features % self.num_heads != 0:
          raise ValueError(
              f"Memory dimension ({self.qkv_features}) must be divisible by "
              f"'num_heads' heads ({self.num_heads})."
          )
  
      self.head_dim = self.qkv_features // self.num_heads
  
      linear_general = functools.partial(
          LinearGeneral,
          in_features=self.in_features,
          out_features=(self.num_heads, self.head_dim),
          dtype=self.dtype,
          param_dtype=self.param_dtype,
          kernel_init=self.kernel_init,
          bias_init=self.bias_init,
          use_bias=self.use_bias,
          precision=self.precision,
          dot_general=self.qkv_dot_general,
          dot_general_cls=self.qkv_dot_general_cls,
      )
      # project inputs_q to multi-headed q/k/v
      # dimensions are then [batch..., length, n_heads, n_features_per_head]
      self.query = linear_general(rngs=rngs)
      self.key = linear_general(rngs=rngs)
      self.value = linear_general(rngs=rngs)
  
      if self.normalize_qk:
          # Normalizing query and key projections stabilizes training with higher
          # LR. See ViT-22B paper http://arxiv.org/abs/2302.05442 for analysis.
          self.query_ln = LayerNorm(
              self.head_dim, use_bias=False, rngs=rngs
          )
          self.key_ln = LayerNorm(
              self.head_dim, use_bias=False, rngs=rngs
          )
      else:
          self.query_ln = None
          self.key_ln = None
  
      self.out = LinearGeneral(
          in_features=(self.num_heads, self.head_dim),
          out_features=self.out_features,
          axis=(-2, -1),
          kernel_init=self.kernel_init,
          bias_init=self.bias_init,
          use_bias=self.use_bias,
          dtype=self.dtype,
          param_dtype=self.param_dtype,
          precision=self.precision,
          dot_general=self.out_dot_general,
          dot_general_cls=self.out_dot_general_cls,
          rngs=rngs,
      )
  
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 1 to create a new MultiHeadAttention object from inherited base classes: Module with example: obj = MultiHeadAttention(num_heads, in_features, qkv_features, out_features)