class_docstr: Multi-head dot-product attention.
  
  Attributes:
    num_heads: number of attention heads. Features (i.e. inputs_q.shape[-1])
      should be divisible by the number of heads.
    head_dim: dimension of each head.
    dtype: the dtype of the computation.
    dropout_rate: dropout rate
    kernel_init: initializer for the kernel of the Dense layers.
    float32_logits: bool, if True then compute logits in float32 to avoid
      numerical issues with bfloat16.
docstr: Applies multi-head dot product attention on the input data.
  
  Projects the inputs into multi-headed query, key, and value vectors,
  applies dot-product attention and project the results to an output vector.
  
  There are two modes: decoding and non-decoding (e.g., training). The mode is
  determined by `decode` argument. For decoding, this method is called twice,
  first to initialize the cache and then for an actual decoding process. The
  two calls are differentiated by the presence of 'cached_key' in the variable
  dict. In the cache initialization stage, the cache variables are initialized
  as zeros and will be filled in the subsequent decoding process.
  
  In the cache initialization call, `inputs_q` has a shape [batch, length,
  q_features] and `inputs_kv`: [batch, length, kv_features]. During the
  incremental decoding stage, query, key and value all have the shape [batch,
  1, qkv_features] corresponding to a single step.
  
  Args:
    inputs_q: input queries of shape `[batch, q_length, q_features]`.
    inputs_kv: key/values of shape `[batch, kv_length, kv_features]`.
    mask: attention mask of shape `[batch, num_heads, q_length, kv_length]`.
    bias: attention bias of shape `[batch, num_heads, q_length, kv_length]`.
    decode: Whether to prepare and use an autoregressive cache.
    deterministic: Disables dropout if set to True.
  
  Returns:
    output of shape `[batch, length, q_features]`.
code: def __call__(
      self,
      inputs_q: constants.JTensor,
      inputs_kv: constants.JTensor,
      mask: Optional[constants.JTensor] = None,
      bias: Optional[constants.JTensor] = None,
      *,
      decode: bool = False,
      deterministic: bool = False,
  ) -> constants.JTensor:
      """Applies multi-head dot product attention on the input data.
  
      Projects the inputs into multi-headed query, key, and value vectors,
      applies dot-product attention and project the results to an output vector.
  
      There are two modes: decoding and non-decoding (e.g., training). The mode is
      determined by `decode` argument. For decoding, this method is called twice,
      first to initialize the cache and then for an actual decoding process. The
      two calls are differentiated by the presence of 'cached_key' in the variable
      dict. In the cache initialization stage, the cache variables are initialized
      as zeros and will be filled in the subsequent decoding process.
  
      In the cache initialization call, `inputs_q` has a shape [batch, length,
      q_features] and `inputs_kv`: [batch, length, kv_features]. During the
      incremental decoding stage, query, key and value all have the shape [batch,
      1, qkv_features] corresponding to a single step.
  
      Args:
        inputs_q: input queries of shape `[batch, q_length, q_features]`.
        inputs_kv: key/values of shape `[batch, kv_length, kv_features]`.
        mask: attention mask of shape `[batch, num_heads, q_length, kv_length]`.
        bias: attention bias of shape `[batch, num_heads, q_length, kv_length]`.
        decode: Whether to prepare and use an autoregressive cache.
        deterministic: Disables dropout if set to True.
  
      Returns:
        output of shape `[batch, length, q_features]`.
      """
      projection = functools.partial(
          t5_layers.DenseGeneral,
          axis=-1,
          features=(self.num_heads, self.head_dim),
          kernel_axes=("embed", "joined_kv"),
          dtype=self.dtype,
      )
  
      # NOTE: T5 does not explicitly rescale the attention logits by
      #       1/sqrt(depth_kq)!  This is folded into the initializers of the
      #       linear transformations, which is equivalent under Adafactor.
      depth_scaling = jnp.sqrt(self.head_dim).astype(self.dtype)
      query_init = lambda *args: self.kernel_init(*args) / depth_scaling
  
      # Project inputs_q to multi-headed q/k/v
      # dimensions are then [batch, length, num_heads, head_dim]
      query = projection(kernel_init=query_init)(inputs_q)
      key = projection(kernel_init=self.kernel_init)(inputs_kv)
      value = projection(kernel_init=self.kernel_init)(inputs_kv)
  
      query = t5_layers.with_sharding_constraint(
          query, ("batch", "length", "heads", "kv")
      )
      key = t5_layers.with_sharding_constraint(
          key, ("batch", "length", "heads", "kv")
      )
      value = t5_layers.with_sharding_constraint(
          value, ("batch", "length", "heads", "kv")
      )
  
      if decode:
          # Detect if we're initializing by absence of existing cache data.
          is_initialized = self.has_variable("cache", "cached_key")
          # The key and value have dimension [batch, length, num_heads, head_dim],
          # but we cache them as [batch, num_heads, head_dim, length] as a TPU
          # fusion optimization. This also enables the "scatter via one-hot
          # broadcast" trick, which means we do a one-hot broadcast instead of a
          # scatter/gather operations, resulting in a 3-4x speedup in practice.
          swap_dims = lambda x: x[:-3] + tuple(
              x[i] for i in [-2, -1, -3]
          )
          cached_key = self.variable(
              "cache",
              "cached_key",
              jnp.zeros,
              swap_dims(key.shape),
              key.dtype,
          )
          cached_value = self.variable(
              "cache",
              "cached_value",
              jnp.zeros,
              swap_dims(value.shape),
              value.dtype,
          )
          cache_index = self.variable(
              "cache",
              "cache_index",
              lambda: jnp.array(0, dtype=jnp.int32),
          )
          if is_initialized:
              (
                  batch,
                  num_heads,
                  head_dim,
                  length,
              ) = cached_key.value.shape
              # During fast autoregressive decoding, we feed one position at a time,
              # and cache the keys and values step by step.
              # Sanity shape check of cached key against input query.
              expected_shape = (batch, 1, num_heads, head_dim)
              if expected_shape != query.shape:
                  raise ValueError(
                      "Autoregressive cache shape error, "
                      "expected query shape %s instead got %s."
                      % (expected_shape, query.shape)
                  )
  
              # Create a OHE of the current index. NOTE: the index is increased below.
              cur_index = cache_index.value
              one_hot_indices = jax.nn.one_hot(
                  cur_index, length, dtype=key.dtype
              )
              # In order to update the key, value caches with the current key and
              # value, we move the length axis to the back, similar to what we did for
              # the cached ones above.
              # Note these are currently the key and value of a single position, since
              # we feed one position at a time.
              one_token_key = jnp.moveaxis(key, -3, -1)
              one_token_value = jnp.moveaxis(value, -3, -1)
              # Update key, value caches with our new 1d spatial slices.
              # We implement an efficient scatter into the cache via one-hot
              # broadcast and addition.
              key = cached_key.value + one_token_key * one_hot_indices
              value = (
                  cached_value.value + one_token_value * one_hot_indices
              )
              cached_key.value = key
              cached_value.value = value
              cache_index.value = cache_index.value + 1
              # Move the keys and values back to their original shapes.
              key = jnp.moveaxis(key, -1, -3)
              value = jnp.moveaxis(value, -1, -3)
  
              # Causal mask for cached decoder self-attention: our single query
              # position should only attend to those key positions that have already
              # been generated and cached, not the remaining zero elements.
              mask = t5_layers.combine_masks(
                  mask,
                  jnp.broadcast_to(
                      jnp.arange(length) <= cur_index,
                      # (1, 1, length) represent (head dim, query length, key length)
                      # query length is 1 because during decoding we deal with one
                      # index.
                      # The same mask is applied to all batch elements and heads.
                      (batch, 1, 1, length),
                  ),
              )
  
              # Grab the correct relative attention bias during decoding. This is
              # only required during single step decoding.
              if bias is not None:
                  # The bias is a full attention matrix, but during decoding we only
                  # have to take a slice of it.
                  # This is equivalent to bias[..., cur_index:cur_index+1, :].
                  bias = t5_layers.dynamic_vector_slice_in_dim(
                      jnp.squeeze(bias, axis=0),
                      jnp.reshape(cur_index, (-1)),
                      1,
                      -2,
                  )
  
      # Convert the boolean attention mask to an attention bias.
      if mask is not None:
          # attention mask in the form of attention bias
          attention_bias = jax.lax.select(
              mask > 0,
              jnp.full(mask.shape, 0.0).astype(self.dtype),
              jnp.full(mask.shape, -1e10).astype(self.dtype),
          )
      else:
          attention_bias = None
  
      # Add provided bias term (e.g. relative position embedding).
      if bias is not None:
          attention_bias = t5_layers.combine_biases(
              attention_bias, bias
          )
  
      dropout_rng = None
      if not deterministic and self.dropout_rate > 0.0:
          dropout_rng = self.make_rng("dropout")
  
      # Apply attention.
      x, attn_weights = dot_product_attention(
          query,
          key,
          value,
          bias=attention_bias,
          dropout_rng=dropout_rng,
          dropout_rate=self.dropout_rate,
          deterministic=deterministic,
          dtype=self.dtype,
          float32_logits=self.float32_logits,
      )
  
      # Back to the original inputs dimensions.
      out = t5_layers.DenseGeneral(
          features=inputs_q.shape[
              -1
          ],  # output dim is set to the input dim.
          axis=(-2, -1),
          kernel_init=self.kernel_init,
          kernel_axes=("joined_kv", "embed"),
          dtype=self.dtype,
      )(x)
      return out, 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 1 for calling the MultiHeadDotProductAttention obj's underlying member method code with example usage: obj.__call__(inputs_q, inputs_kv, mask, bias) and returns: out,  attn_weights