code: def multi_head_dot_product_attention(
      scope: Scope,
      inputs_q,
      inputs_kv,
      num_heads,
      dtype=jnp.float32,
      qkv_features=None,
      out_features=None,
      attention_axis=None,
      causal_mask=False,
      padding_mask=None,
      key_padding_mask=None,
      segmentation=None,
      key_segmentation=None,
      cache=False,
      broadcast_dropout=True,
      dropout_rng=None,
      dropout_rate=0.0,
      deterministic=False,
      precision=None,
      kernel_init=default_kernel_init,
      bias_init=initializers.zeros_init(),
      bias=True,
      attention_fn=dot_product_attention,
  ):
      """Applies multi-head dot product attention on the input data.
  
      Projects the inputs into multi-headed query, key, and value vectors,
      applies dot-product attention and project the results to an output vector.
  
      This can be used for encoder-decoder attention by specifying both `inputs_q`
      and `inputs_kv` orfor self-attention by only specifying `inputs_q` and
      setting `inputs_kv` to None.
  
      Args:
        inputs_q: input queries of shape `[bs, dim1, dim2, ..., dimN, features]`.
        inputs_kv: key/values of shape `[bs, dim1, dim2, ..., dimN, features]`
          or None for self-attention, inn which case key/values will be derived
          from inputs_q.
        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: float32)
        qkv_features: dimension of the key, query, and value.
        out_features: dimension of the last projection
        attention_axis: axes over which the attention is applied ( 'None' means
          attention over all axes, but batch, heads, and features).
        causal_mask: boolean specifying whether to apply a causal mask on the
          attention weights. If True, the output at timestep `t` will not depend
          on inputs at timesteps strictly greater than `t`.
        padding_mask: boolean specifying query tokens that are pad token.
        key_padding_mask: boolean specifying key-value tokens that are pad token.
        segmentation: segment indices for packed inputs_q data.
        key_segmentation: segment indices for packed inputs_kv data.
        cache: an instance of `flax.deprecated.nn.attention.Cache` used for
          efficient autoregressive decoding.
        broadcast_dropout: bool: use a broadcasted dropout along batch dims.
        dropout_rng: JAX PRNGKey: to be used for dropout
        dropout_rate: dropout rate
        deterministic: bool, deterministic or not (to apply dropout)
        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.
        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]``
  
      Returns:
        output of shape `[bs, dim1, dim2, ..., dimN, features]`.
      """
  
      assert (
          causal_mask or not cache
      ), "Caching is only support for causal attention."
  
      if inputs_kv is None:
          inputs_kv = inputs_q
  
      if attention_axis is None:
          attention_axis = tuple(range(1, inputs_q.ndim - 1))
  
      features = out_features or inputs_q.shape[-1]
      qkv_features = qkv_features or inputs_q.shape[-1]
  
      assert (
          qkv_features % num_heads == 0
      ), "Memory dimension must be divisible by number of heads."
      head_dim = qkv_features // num_heads
  
      dense = functools.partial(
          dense_general,
          axis=-1,
          dtype=dtype,
          features=(num_heads, head_dim),
          kernel_init=kernel_init,
          bias_init=bias_init,
          bias=bias,
          precision=precision,
      )
      # project inputs_q to multi-headed q/k/v
      # dimensions are then [bs, dims..., n_heads, n_features_per_head]
      query = scope.child(dense, "query")(inputs_q)
      key = scope.child(dense, "key")(inputs_kv)
      value = scope.child(dense, "value")(inputs_kv)
  
      if cache:
          cache_entry: Union[Callable[[Any], CacheEntry], CacheEntry]
          if not scope.has_variable("cache", "entry"):
              ndim, tail_shape = (key.ndim, key.shape[-2:])
  
              def init_fn(shape, dtype=jnp.float32):
                  full_shape = shape + tail_shape
                  if len(full_shape) != ndim:
                      raise ValueError(
                          "Shape should be a tuple with the shape of the batch"
                          "and attention dims."
                      )
                  return CacheEntry(
                      key=jnp.zeros(full_shape, dtype),
                      value=jnp.zeros(full_shape, dtype),
                      i=jnp.zeros((), jnp.uint32),
                  )
  
              cache_entry = init_fn
          else:
              cache_entry = scope.get_variable("cache", "entry")
              if not isinstance(cache_entry, CacheEntry):
                  raise ValueError("Cache is not initialized.")
  
              expected_shape = list(cache_entry.key.shape[:-2])
              for attn_dim in attention_axis:
                  expected_shape[attn_dim] = 1
              expected_shape = (
                  tuple(expected_shape) + inputs_q.shape[-1:]
              )
              if expected_shape != inputs_q.shape:
                  raise ValueError(
                      "Invalid shape provided, expected shape %s instead got %s."
                      % (expected_shape, inputs_q.shape)
                  )
  
              cshape = cache_entry.key.shape
              indices = [0] * len(cshape)
              i = cache_entry.i
              attn_size = np.prod(np.take(cshape, attention_axis))
              for attn_dim in attention_axis:
                  attn_size //= cshape[attn_dim]
                  indices[attn_dim] = i // attn_size
                  i = i % attn_size
  
              key = lax.dynamic_update_slice(cache_entry.key, key, indices)  # type: ignore
              value = lax.dynamic_update_slice(cache_entry.value, value, indices)  # type: ignore
              one = jnp.array(1, jnp.uint32)
              cache_entry = cache_entry.replace(
                  i=cache_entry.i + one, key=key, value=value
              )
  
              # TODO(levskaya): verify this is still needed in translation decoding.
              key_padding_mask = jnp.broadcast_to(
                  (jnp.arange(cshape[1]) < cache_entry.i), cshape[:2]
              )
              key_padding_mask = key_padding_mask.astype(jnp.float32)[
                  ..., None
              ]
          scope.put_variable("cache", "entry", cache_entry)
  
      # create attention masks
      mask_components = []
  
      if causal_mask:
          if cache and isinstance(cache_entry, CacheEntry):
              bias_pre_shape = (1,) * (key.ndim - 1)
              attn_shape = tuple(np.take(key.shape, attention_axis))
              attn_size = np.prod(attn_shape)
              ii = jnp.arange(attn_size, dtype=jnp.uint32)
              mask = ii < cache_entry.i
              mask_components.append(
                  mask.reshape(bias_pre_shape + attn_shape)
              )
          else:
              mask_components.append(
                  _make_causal_mask(key, attention_axis)
              )
  
      if padding_mask is not None:
          if key_padding_mask is None:
              key_padding_mask = padding_mask
          padding_mask = make_padding_mask(
              padding_mask_query=padding_mask,
              padding_mask_key=key_padding_mask,
              query_shape=query.shape,
              key_shape=key.shape,
              attention_axis=attention_axis,
          )
          mask_components.append(padding_mask)
  
      if segmentation is not None:
          if key_segmentation is None:
              key_segmentation = segmentation
          segmentation_mask = make_padding_mask(
              padding_mask_query=segmentation,
              padding_mask_key=key_segmentation,
              query_shape=query.shape,
              key_shape=key.shape,
              attention_axis=attention_axis,
              segmentation_mask=True,
          )
          mask_components.append(segmentation_mask)
  
      if mask_components:
          attention_mask = mask_components[0]
          for component in mask_components[1:]:
              attention_mask = jnp.logical_and(
                  attention_mask, component
              )
  
          # attention mask in the form of attention bias
          attention_bias = lax.select(
              attention_mask > 0,
              jnp.full(attention_mask.shape, 0.0).astype(dtype),
              jnp.full(attention_mask.shape, -1e10).astype(dtype),
          )
      else:
          attention_bias = None
  
      # apply attention
      x = scope.child(attention_fn)(
          query,
          key,
          value,
          dtype=dtype,
          axis=attention_axis,
          bias=attention_bias,
          precision=precision,
          dropout_rng=dropout_rng,
          dropout_rate=dropout_rate,
          broadcast_dropout=broadcast_dropout,
          deterministic=deterministic,
      )
  
      # back to the original inputs dimensions
      out = scope.child(dense_general, name="out")(
          x,
          features=features,
          axis=(-2, -1),
          kernel_init=kernel_init,
          bias_init=bias_init,
          bias=bias,
          dtype=dtype,
          precision=precision,
      )
  
      return out
  
file_path: flax/flax/core/nn/attention.py
imports: import functools
  from collections.abc import Iterable
  from typing import Any, Callable, Union
  import jax
  import jax.numpy
  import numpy
  from jax import lax, random
  from flax import struct
  from flax.core import Scope
  from flax.linen import initializers
  from .linear import default_kernel_init, dense_general
label_desc: Use image node_id 3 for calling a global function with example usage: multi_head_dot_product_attention(scope, inputs_q, inputs_kv, num_heads, dtype, qkv_features, out_features, attention_axis, causal_mask, padding_mask, key_padding_mask, segmentation, key_segmentation, cache, broadcast_dropout, dropout_rng, dropout_rate, deterministic, precision, kernel_init, bias_init, bias, attention_fn) and returns: out, CacheEntry