code: def axial_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,
  ) -> jnp.ndarray:
      """Applies masked, head-axial qkv dot-product attention.
  
      Assigns different heads for different axes which is more efficient and
      allows for having attention on all axes in every layer.
  
      Args:
        query: Queries for calculating attention with shape of `[batch, ...,
          num_heads, qk_depth_per_head]`.
        key: Keys for calculating attention with shape of `[batch, ..., num_heads,
          qk_depth_per_head]`.
        value: Values to be used in attention with shape of `[batch, ..., num_heads,
          v_depth_per_head]`.
        bias: Bias is not supported and will raise an error if passed.
        bias_kv: Bias for the attention weights. This should be
          broadcastable to the shape: `[batch, ...]`. This can be used for
            incorporating causal masks, padding masks, proximity bias, etc. Default
            is None, which means no bias is applied on attention matrix.
        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.
  
      Returns:
        Output of shape `[bs, ..., num_heads, features]`.
      """
      if query.shape != key.shape:
          raise ValueError(
              "Axial dot product attention only supports "
              "query and key with the same shape."
          )
      if bias is not None:
          raise NotImplementedError("Bias passed to axial attention.")
      if bias_kv is not None:
          # expand padding mask for head dimension, and last dimension, which will
          # be broadcasted. [batch, ..., 1, 1]
          bias_kv = bias_kv[..., jnp.newaxis, jnp.newaxis]
      # Normalize the query with the squre of its depth.
      query = query / jnp.sqrt(query.shape[-1]).astype(dtype)
      prefix_str = "abcdefghijk"
      # Split heads for each axial attention dimension.
      num_attn_dimensions = (
          query.ndim - 3
      )  # all dims but bs, heads, and channel.
      if query.shape[-2] % num_attn_dimensions != 0:
          raise ValueError(
              f"In head-axial dot-product attention, number of "
              f"heads ({query.shape[-2]}) should be divisible by number "
              f"of attention dimensions ({num_attn_dimensions})!"
          )
  
      queries = jnp.split(query, num_attn_dimensions, axis=-2)
      keys = jnp.split(key, num_attn_dimensions, axis=-2)
      values = jnp.split(value, num_attn_dimensions, axis=-2)
  
      outputs = []
      for i, (query, key, value) in enumerate(
          zip(queries, keys, values)
      ):
          axis = i + 1  # + 1 for batch
          batch_dims = prefix_str[:axis]
          einsum_str = (
              f"{batch_dims}x...z,{batch_dims}y...z->{batch_dims}x...y"
          )
          attn_logits = jnp.einsum(
              einsum_str, query, key, precision=precision
          )
          if bias_kv is not None:
              # put attention axis into last dimension
              attn_logits += jnp.swapaxes(
                  bias_kv, axis, -1
              )  # {batch_dims}1...y
          attn_weights = jax.nn.softmax(attn_logits, axis=-1)
  
          # Apply dropout.
          if not deterministic and dropout_rate > 0.0:
              attn_weights = _attention_dropout(
                  attn_weights,
                  rate=dropout_rate,
                  broadcast=broadcast_dropout,
                  dropout_rng=dropout_rng,
              )
          einsum_str = (
              f"{batch_dims}x...y,{batch_dims}y...z->{batch_dims}x...z"
          )
          outputs.append(
              jnp.einsum(
                  einsum_str, attn_weights, value, precision=precision
              )
          )
  
      return jnp.concatenate(outputs, axis=-2)
  
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 3 for calling a global function with example usage: axial_dot_product_attention(query, key, value) and returns: jnp