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[Array] = 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
          ``nnx.Intermediate`` collection. 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: 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 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