code: def dot_product_attention(
      query: constants.JTensor,
      key: constants.JTensor,
      value: constants.JTensor,
      bias: Optional[constants.JTensor] = None,
      dropout_rng: Optional[constants.JTensor] = None,
      dropout_rate: float = 0.0,
      deterministic: bool = False,
      dtype: constants.DType = jnp.float32,
      float32_logits: bool = False,
  ):
      """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.
  
      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.
        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: float32)
        float32_logits: bool, if True then compute logits in float32 to avoid
          numerical issues with bfloat16.
  
      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."
  
      # Casting logits and softmax computation for float32 for model stability.
      if float32_logits:
          query = query.astype(jnp.float32)
          key = key.astype(jnp.float32)
  
      # `attn_weights`: [batch, num_heads, q_length, kv_length]
      attn_weights = jnp.einsum("bqhd,bkhd->bhqk", query, key)
  
      # Apply attention bias: masking, dropout, proximity bias, etc.
      if bias is not None:
          attn_weights = attn_weights + bias.astype(attn_weights.dtype)
  
      # Normalize the attention weights across `kv_length` dimension.
      attn_weights = jax.nn.softmax(attn_weights).astype(dtype)
  
      # Apply attention dropout.
      if not deterministic and dropout_rate > 0.0:
          keep_prob = 1.0 - dropout_rate
          # T5 broadcasts along the "length" dim, but unclear which one that
          # corresponds to in positional dimensions here, assuming query dim.
          dropout_shape = list(attn_weights.shape)
          dropout_shape[-2] = 1
          keep = jax.random.bernoulli(
              dropout_rng, keep_prob, dropout_shape
          )
          keep = jnp.broadcast_to(keep, attn_weights.shape)
          multiplier = keep.astype(attn_weights.dtype) / jnp.asarray(
              keep_prob, dtype=dtype
          )
          attn_weights = attn_weights * multiplier
  
      # Take the linear combination of `value`.
      return (
          jnp.einsum("bhqk,bkhd->bqhd", attn_weights, value),
          attn_weights,
      )
  
imports: import functools
  from typing import Optional, Sequence
  from flax import linen
  import jax
  import jax.numpy
  from scenic.projects.knowledge_visual_language.models import constants
  from t5x.examples.t5 import layers
  from t5x.examples.t5 import network
label_desc: Use image node_id 5 for calling a global function with example usage: dot_product_attention(query, key, value, bias, dropout_rng, dropout_rate, deterministic, dtype, float32_logits) and returns: jnp,  attn_weights