code: def dot_product_attention(
      query: Array,
      key: Array,
      value: Array,
      bias: Optional[Array] = None,
      mask: Optional[Array] = None,
      broadcast_dropout: bool = True,
      dropout_rng: Optional[PRNGKey] = None,
      dropout_rate: float = 0.0,
      deterministic: bool = False,
      dtype: Optional[Dtype] = None,
      precision: PrecisionLike = None,
      module: Optional[Module] = 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.
  
      Note: query, key, value needn't have any batch dimensions.
  
      Args:
        query: queries for calculating attention with shape of ``[batch..., q_length,
          num_heads, qk_depth_per_head]``.
        key: keys for calculating attention with shape of ``[batch..., kv_length,
          num_heads, qk_depth_per_head]``.
        value: values to be used in attention with shape of ``[batch..., kv_length,
          num_heads, v_depth_per_head]``.
        bias: bias for the attention weights. This should be broadcastable to the
          shape ``[batch..., num_heads, q_length, kv_length]``. This can be used for
          incorporating causal masks, padding masks, proximity bias, etc.
        mask: mask for the attention weights. This should be broadcastable to the
          shape ``[batch..., num_heads, q_length, kv_length]``. This can be used for
          incorporating causal masks. Attention weights are masked out if their
          corresponding mask value is ``False``.
        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)
        dtype: the dtype of the computation (default: infer from inputs)
        precision: numerical precision of the computation see ``jax.lax.Precision`
          for details.
        module: the Module that will sow the attention weights into the
          'intermediates' collection. Remember to mark 'intermediates' as mutable via
          ``mutable=['intermediates']`` in order to have that collection returned.
          If ``module`` is None, the attention weights will not be sowed.
  
      Returns:
        Output of shape ``[batch..., q_length, num_heads, v_depth_per_head]``.
      """
      query, key, value = promote_dtype(query, key, value, dtype=dtype)
      dtype = query.dtype
      assert (
          key.ndim == query.ndim == value.ndim
      ), "q, k, v must have same rank."
      assert (
          query.shape[:-3] == key.shape[:-3] == value.shape[:-3]
      ), "q, k, v batch dims must match."
      assert (
          query.shape[-2] == key.shape[-2] == value.shape[-2]
      ), "q, k, v num_heads must match."
      assert (
          key.shape[-3] == value.shape[-3]
      ), "k, v lengths must match."
  
      # compute attention weights
      attn_weights = dot_product_attention_weights(
          query,
          key,
          bias,
          mask,
          broadcast_dropout,
          dropout_rng,
          dropout_rate,
          deterministic,
          dtype,
          precision,
          module,
      )
  
      # return weighted sum over values for each query position
      return jnp.einsum(
          "...hqk,...khd->...qhd",
          attn_weights,
          value,
          precision=precision,
      )
  
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 a global function with example usage: dot_product_attention(query, key, value, bias, mask, broadcast_dropout, dropout_rng, dropout_rate, deterministic, dtype, precision, module) and returns: jnp