code: def dot_product_attention(
      query: jnp.ndarray,
      key: jnp.ndarray,
      value: jnp.ndarray,
      *,
      bias: Optional[jnp.ndarray] = None,
      bias_kv: Optional[jnp.ndarray] = None,
      broadcast_dropout: bool = True,
      dropout_rate: float = 0.1,
      dtype: jnp.dtype = jnp.float32,
      precision: Optional[jax.lax.Precision] = None,
      deterministic: bool,
      dropout_rng: Optional[jnp.ndarray] = None,
      capture_attention_weights: bool = True,
  ) -> jnp.ndarray:
      """Computes the 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.
        bias_kv: Attention bias defined for keys only which has shape
          `[batch..., kv_length]`. Can be used for masking elements in k/v.
        broadcast_dropout: Use a broadcasted dropout along batch dims.
        dropout_rate: Dropout rate.
        dtype: The dtype of the computation (default: float32).
        precision: Numerical precision of the computation see `jax.lax.Precision`
          for details.
        deterministic: Deterministic or not (to apply dropout).
        dropout_rng: Optional JAX PRNGKey to be used for dropout.
        capture_attention_weights: Whether to add an identity layer to tag the
          attention weights to be used for capturing them using Flax
          capture_intermediate, e.g. for visualization. Note that if this is set to
          True, this function can be only called within a Flax module.
  
      Returns:
        Output of shape `[batch..., length, num_heads, v_depth_per_head]`.
      """
      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."
      assert query.shape[-1] == key.shape[-1], "q, k depths must match."
  
      # Calculate attention matrix.
      depth = query.shape[-1]
      query = query / jnp.sqrt(depth).astype(dtype)
      # attn weight shape is (batch..., num_heads, q_length, kv_length)
      attn_weights = jnp.einsum(
          "...qhd,...khd->...hqk", query, key, precision=precision
      )
  
      # Apply attention bias: masking, dropout, proximity bias, etc.
      if bias is not None:
          attn_weights = attn_weights + bias
      if bias_kv is not None:
          bias_kv = bias_kv[..., jnp.newaxis, jnp.newaxis, :]
          attn_weights += bias_kv
  
      # Normalize the attention weights.
      attn_weights = jax.nn.softmax(attn_weights).astype(dtype)
  
      if capture_attention_weights:
          # Tag the intermediate weights for logging/visualization.
          attn_weights = nn_layers.IdentityLayer(name="attn_weights")(
              attn_weights
          )
  
      # Apply attention dropout.
      if not deterministic and dropout_rate > 0.0:
          if dropout_rng is None:
              raise ValueError(
                  "Did not provide `rng` to dot_product_attention()."
              )
          attn_weights = _attention_dropout(
              attn_weights,
              rate=dropout_rate,
              broadcast=broadcast_dropout,
              dropout_rng=dropout_rng,
          )
  
      # Return weighted sum over values for each query position.
      return jnp.einsum(
          "...hqk,...khd->...qhd",
          attn_weights,
          value,
          precision=precision,
      )
  
imports: import functools
  from typing import Callable, Optional, Sequence, Tuple
  import flax.linen
  import jax
  import jax.numpy
  import numpy
  from scenic.model_lib.layers import nn_layers
label_desc: Use image node_id 2 for calling a global function with example usage: dot_product_attention(query, key, value) and returns: jnp