class_docstr: Customized dynamic 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.
    use_bias: Whether pointwise QKV dense transforms use bias.
    precision: Numerical precision of the computation see `jax.lax.Precision`
      for details.
    dtype: the dtype of the computation (default: float32).
class_docstr_tok: ["Customized","dynamic","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",".","use_bias",":","Whether","pointwise","QKV","dense","transforms","use","bias",".","precision",":","Numerical","precision","of","the","computation","see","`","jax.lax.Precision","`","for","details",".","dtype",":","the","dtype","of","the","computation","(","default",":","float32",")","."]
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, inn which case key/values will be derived from inputs_q.
    dataset: Current dataset.
    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",",","inn","which","case","key\/values","will","be","derived","from","inputs_q",".","dataset",":","Current","dataset",".","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],
      dataset: str,
      *,
      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, inn which case key/values will be derived from inputs_q.
        dataset: Current dataset.
        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]
  
      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")
  
      attention_fn = attention_layers.dot_product_attention
      x = 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 = nn.DenseGeneral(
          features=features,
          axis=(-2, -1),
          kernel_init=self.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 Any, Callable, Optional, Sequence, Tuple, Union
  import flax.linen
  import jax
  from jax.nn import initializers
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.model_lib.layers import attention_layers
  from scenic.model_lib.layers import nn_layers
  from scenic.projects.baselines import vit
  from scenic.projects.polyvit import polyvit_base_model
  from scenic.projects.vivit import model
label_desc: Use image node_id 1 for calling the DynamicMultiHeadAttention obj's underlying member method code with example usage: obj.__call__(inputs_q, inputs_kv, dataset) and returns: out, unknown