code: def dot_product_attention(
      scope,
      query,
      key,
      value,
      dtype=jnp.float32,
      bias=None,
      axis=None,
      broadcast_dropout=True,
      dropout_rng=None,
      dropout_rate=0.0,
      deterministic=False,
      precision=None,
  ):
      """Computes dot-product attention given query, key, and value.
  
      This is the core function for applying attention based on
      https://arxiv.org/abs/1706.03762. It calculates the attention weights given
      query and key and combines the values using the attention weights. This
      function supports multi-dimensional inputs.
  
  
      Args:
        query: queries for calculating attention with shape of `[batch_size, dim1,
          dim2, ..., dimN, num_heads, mem_channels]`.
        key: keys for calculating attention with shape of `[batch_size, dim1, dim2,
          ..., dimN, num_heads, mem_channels]`.
        value: values to be used in attention with shape of `[batch_size, dim1,
          dim2,..., dimN, num_heads, value_channels]`.
        dtype: the dtype of the computation (default: float32)
        bias: bias for the attention weights. This can be used for incorporating
          autoregressive mask, padding mask, proximity bias.
        axis: axises over which the attention is applied.
        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.
  
      Returns:
        Output of shape `[bs, dim1, dim2, ..., dimN,, num_heads, value_channels]`.
      """
      assert key.shape[:-1] == value.shape[:-1]
      assert (
          query.shape[0:1] == key.shape[0:1]
          and query.shape[-1] == key.shape[-1]
      )
  
      if axis is None:
          axis = tuple(range(1, key.ndim - 2))
      if not isinstance(axis, Iterable):
          axis = (axis,)
      assert key.ndim == query.ndim
      assert key.ndim == value.ndim
      for ax in axis:
          if not (query.ndim >= 3 and 1 <= ax < query.ndim - 2):
              raise ValueError(
                  "Attention axis must be between the batch axis and the last-two axes."
              )
      depth = query.shape[-1]
      n = key.ndim
      # batch_dims is  <bs, <non-attention dims>, num_heads>
      batch_dims = tuple(np.delete(range(n), axis + (n - 1,)))
      # q & k -> (bs, <non-attention dims>, num_heads, <attention dims>, channels)
      qk_perm = batch_dims + axis + (n - 1,)
      key = key.transpose(qk_perm)
      query = query.transpose(qk_perm)
      # v -> (bs, <non-attention dims>, num_heads, channels, <attention dims>)
      v_perm = batch_dims + (n - 1,) + axis
      value = value.transpose(v_perm)
  
      query = query / jnp.sqrt(depth).astype(dtype)
      batch_dims_t = tuple(range(len(batch_dims)))
      attn_weights = lax.dot_general(
          query,
          key,
          (((n - 1,), (n - 1,)), (batch_dims_t, batch_dims_t)),
          precision=precision,
      )
  
      # apply attention bias: masking, droput, proximity bias, ect.
      if bias is not None:
          attn_weights = attn_weights + bias
  
      # normalize the attention weights
      norm_dims = tuple(
          range(attn_weights.ndim - len(axis), attn_weights.ndim)
      )
      attn_weights = lax.exp(
          attn_weights
          - jax.scipy.special.logsumexp(
              attn_weights, axis=norm_dims, keepdims=True
          )
      )
      attn_weights = attn_weights.astype(dtype)
  
      # apply dropout
      if not deterministic and dropout_rate > 0.0:
          if dropout_rng is None:
              dropout_rng = scope.make_rng("dropout")
          keep_prob = jax.lax.tie_in(attn_weights, 1.0 - dropout_rate)
          if broadcast_dropout:
              # dropout is broadcast across the batch+head+non-attention dimension
              dropout_dims = attn_weights.shape[-(2 * len(axis)) :]
              dropout_shape = (
                  tuple([1] * len(batch_dims_t)) + dropout_dims
              )
              keep = random.bernoulli(
                  dropout_rng, keep_prob, dropout_shape
              )
          else:
              keep = random.bernoulli(
                  dropout_rng, keep_prob, attn_weights.shape
              )
          multiplier = keep.astype(attn_weights.dtype) / jnp.asarray(
              keep_prob, dtype=dtype
          )
          attn_weights = attn_weights * multiplier
  
      # compute the new values given the attention weights
      wv_contracting_dims = (
          norm_dims,
          range(value.ndim - len(axis), value.ndim),
      )
      y = lax.dot_general(
          attn_weights,
          value,
          (wv_contracting_dims, (batch_dims_t, batch_dims_t)),
          precision=precision,
      )
  
      # back to (bs, dim1, dim2, ..., dimN, num_heads, channels)
      perm_inv = _invert_perm(qk_perm)
      y = y.transpose(perm_inv)
      return y
  
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 1 for calling a global function with example usage: dot_product_attention(scope, query, key, value, dtype, bias, axis, broadcast_dropout, dropout_rng, dropout_rate, deterministic, precision) and returns: y