class_docstr: Customized multi-head attention for scenic.
  
  Attributes:
    num_heads: Number of attention heads. Features (i.e. inputs_q.shape[-1])
      should be divisible by the number of heads.
    qkv_features: Dimension of the key, query, and value.
    out_features: Dimension of the last projection.
    dropout_rate: Dropout rate.
    broadcast_dropout: Use a broadcasted dropout along batch dims.
    kernel_init: Initializer for the kernel of the Dense layers.
    bias_init: Initializer for the bias of the Dense layers.
    out_kernel_init: Initializer for the kernel of the output Dense layers. If
      None, kernel_init will be used.
    use_bias: Whether pointwise QKV dense transforms use bias.
    precision: Numerical precision of the computation see `jax.lax.Precision`
      for details.
    attention_fn: Defaults to dot_product_attention. Other function of the
      same signature are possible.
    dtype: the dtype of the computation (default: float32).
    enforce_hidden_size_divisible_by_heads: Whether or not we allow the hidden
      size to not be divisible by the number of heads.
class_docstr_tok: ["Customized","multi-head","attention","for","scenic",".","Attributes",":","num_heads",":","Number","of","attention","heads",".","Features","(","i.e",".","inputs_q.shape","[","-1","]",")","should","be","divisible","by","the","number","of","heads",".","qkv_features",":","Dimension","of","the","key",",","query",",","and","value",".","out_features",":","Dimension","of","the","last","projection",".","dropout_rate",":","Dropout","rate",".","broadcast_dropout",":","Use","a","broadcasted","dropout","along","batch","dims",".","kernel_init",":","Initializer","for","the","kernel","of","the","Dense","layers",".","bias_init",":","Initializer","for","the","bias","of","the","Dense","layers",".","out_kernel_init",":","Initializer","for","the","kernel","of","the","output","Dense","layers",".","If","None",",","kernel_init","will","be","used",".","use_bias",":","Whether","pointwise","QKV","dense","transforms","use","bias",".","precision",":","Numerical","precision","of","the","computation","see","`","jax.lax.Precision","`","for","details",".","attention_fn",":","Defaults","to","dot_product_attention",".","Other","function","of","the","same","signature","are","possible",".","dtype",":","the","dtype","of","the","computation","(","default",":","float32",")",".","enforce_hidden_size_divisible_by_heads",":","Whether","or","not","we","allow","the","hidden","size","to","not","be","divisible","by","the","number","of","heads","."]
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.
  
  This can be used for encoder-decoder attention by specifying both `inputs_q`
  and `inputs_kv` or for self-attention by only specifying `inputs_q` and
  setting `inputs_kv` to None.
  
  Args:
    inputs_q: Input queries of shape  `[bs, ..., len_q, features]`.
    inputs_kv: Key/values of shape `[bs, ..., len_k, features]` or None for
      self-attention, in which case key/values will be derived from inputs_q.
    pos_emb_q: Positional embedding to be added to the query.
    pos_emb_k: Positional embedding to be added to the key.
    pos_emb_v: Positional embedding to be added to the value.
    attention_bias: Full attention bias. Should be broadcastable to:
      inputs_q.shape[:-2] + (num_heads, len_q, len_k).
    attention_bias_kv: Attention bias for keys independent of queries which
      has shape (bs, ..., len_k).
    deterministic: Run deterministically or with dropout.
  
  Returns:
    Output of shape `[bs, ..., features]`.
docstr_tok: ["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",".","This","can","be","used","for","encoder-decoder","attention","by","specifying","both","`","inputs_q","`","and","`","inputs_kv","`","or","for","self-attention","by","only","specifying","`","inputs_q","`","and","setting","`","inputs_kv","`","to","None",".","Args",":","inputs_q",":","Input","queries","of","shape","`","[","bs",",","...",",","len_q",",","features","]","`",".","inputs_kv",":","Key\/values","of","shape","`","[","bs",",","...",",","len_k",",","features","]","`","or","None","for","self-attention",",","in","which","case","key\/values","will","be","derived","from","inputs_q",".","pos_emb_q",":","Positional","embedding","to","be","added","to","the","query",".","pos_emb_k",":","Positional","embedding","to","be","added","to","the","key",".","pos_emb_v",":","Positional","embedding","to","be","added","to","the","value",".","attention_bias",":","Full","attention","bias",".","Should","be","broadcastable","to",":","inputs_q.shape","[",":","-2","]","+","(","num_heads",",","len_q",",","len_k",")",".","attention_bias_kv",":","Attention","bias","for","keys","independent","of","queries","which","has","shape","(","bs",",","...",",","len_k",")",".","deterministic",":","Run","deterministically","or","with","dropout",".","Returns",":","Output","of","shape","`","[","bs",",","...",",","features","]","`","."]
code: def __call__(
      self,
      inputs_q: jnp.ndarray,
      inputs_kv: Optional[jnp.ndarray],
      *,
      pos_emb_q: Optional[jnp.ndarray] = None,
      pos_emb_k: Optional[jnp.ndarray] = None,
      pos_emb_v: Optional[jnp.ndarray] = None,
      attention_bias: Optional[jnp.ndarray] = None,
      attention_bias_kv: Optional[jnp.ndarray] = None,
      deterministic: bool = False,
  ) -> jnp.ndarray:
      """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.
  
      This can be used for encoder-decoder attention by specifying both `inputs_q`
      and `inputs_kv` or for self-attention by only specifying `inputs_q` and
      setting `inputs_kv` to None.
  
      Args:
        inputs_q: Input queries of shape  `[bs, ..., len_q, features]`.
        inputs_kv: Key/values of shape `[bs, ..., len_k, features]` or None for
          self-attention, in which case key/values will be derived from inputs_q.
        pos_emb_q: Positional embedding to be added to the query.
        pos_emb_k: Positional embedding to be added to the key.
        pos_emb_v: Positional embedding to be added to the value.
        attention_bias: Full attention bias. Should be broadcastable to:
          inputs_q.shape[:-2] + (num_heads, len_q, len_k).
        attention_bias_kv: Attention bias for keys independent of queries which
          has shape (bs, ..., len_k).
        deterministic: Run deterministically or with dropout.
  
      Returns:
        Output of shape `[bs, ..., features]`.
      """
      if inputs_kv is None:
          inputs_kv = inputs_q
  
      features = self.out_features or inputs_q.shape[-1]
      qkv_features = self.qkv_features or inputs_q.shape[-1]
  
      if self.enforce_hidden_size_divisible_by_heads:
          assert (
              qkv_features % self.num_heads == 0
          ), "Memory dimension must be divisible by number of heads."
      head_dim = qkv_features // self.num_heads
  
      def add_positional_emb(x, pos):
          return x + pos if pos is not None else x
  
      query, key, value = (
          add_positional_emb(inputs_q, pos_emb_q),
          add_positional_emb(inputs_kv, pos_emb_k),
          add_positional_emb(inputs_kv, pos_emb_v),
      )
  
      dense = functools.partial(
          nn.DenseGeneral,
          axis=-1,
          features=(self.num_heads, head_dim),
          kernel_init=self.kernel_init,
          bias_init=self.bias_init,
          use_bias=self.use_bias,
          precision=self.precision,
      )
      # Project inputs_q to multi-headed q/k/v.
      # Dimensions are then [..., l, n_heads, n_features_per_head].
      query, key, value = (
          dense(name="query")(query),
          dense(name="key")(key),
          dense(name="value")(value),
      )
  
      # pylint: disable=too-many-function-args
      attn_kwargs = {}
      if attention_bias_kv is not None:
          # Not necessarily supported by all underlying functions.
          attn_kwargs["bias_kv"] = attention_bias_kv
      if not deterministic and self.dropout_rate > 0:
          attn_kwargs["dropout_rng"] = self.make_rng("dropout")
  
      x = self.attention_fn(
          query,
          key,
          value,
          bias=attention_bias,
          dropout_rate=self.dropout_rate,
          broadcast_dropout=self.broadcast_dropout,
          deterministic=deterministic,
          dtype=self.dtype,
          precision=self.precision,
          **attn_kwargs,
      )
      # pylint: enable=too-many-function-args
  
      # Back to the original inputs dimensions.
      out_kernel_init = (
          self.out_kernel_init
          if self.out_kernel_init is not None
          else self.kernel_init
      )
      out = nn.DenseGeneral(
          features=features,
          axis=(-2, -1),
          kernel_init=out_kernel_init,
          bias_init=self.bias_init,
          use_bias=True,
          dtype=self.dtype,
          precision=self.precision,
          name="out",
      )(x)
  
      return out
  
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 1 for calling the MultiHeadAttention obj's underlying member method code with example usage: obj.__call__(inputs_q, inputs_kv) and returns: out, unknown