class_docstr: DETR Customized 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.
    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.
    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: bool: whether pointwise QKV dense transforms use bias. In DETR
      they always have a bias on the output.
    dtype: the dtype of the computation (default: float32)
class_docstr_tok: ["DETR","Customized","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",".","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",".","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",":","bool",":","whether","pointwise","QKV","dense","transforms","use","bias",".","In","DETR","they","always","have","a","bias","on","the","output",".","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, features]`.
    inputs_kv: Key/values of shape `[bs, len, 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.
    key_padding_mask: Binary array. Key-value tokens that are padded are 0,
      and 1 otherwise.
    train: Train or not (to apply dropout).
  
  Returns:
    output of shape `[bs, len, features]`.
code: def __call__(
      self,
      inputs_q: jnp.ndarray,
      inputs_kv: Optional[jnp.ndarray] = None,
      *,
      pos_emb_q: Optional[jnp.ndarray] = None,
      pos_emb_k: Optional[jnp.ndarray] = None,
      pos_emb_v: Optional[jnp.ndarray] = None,
      key_padding_mask: Optional[jnp.ndarray] = None,
      train: 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, features]`.
        inputs_kv: Key/values of shape `[bs, len, 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.
        key_padding_mask: Binary array. Key-value tokens that are padded are 0,
          and 1 otherwise.
        train: Train or not (to apply dropout).
  
      Returns:
        output of shape `[bs, len, features]`.
      """
      if inputs_kv is None:
          inputs_kv = inputs_q
  
      assert inputs_kv.ndim == inputs_q.ndim == 3
      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_emb_x):
          return x if pos_emb_x is None else x + pos_emb_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,
          dtype=self.dtype,
      )
      # project inputs_q to multi-headed q/k/v
      # dimensions are then [bs, l, n_heads, n_features_per_head]
      query, key, value = (
          dense(name="query")(query),
          dense(name="key")(key),
          dense(name="value")(value),
      )
  
      # create attention masks
      if key_padding_mask is not None:
          attention_bias = (1 - key_padding_mask) * -1e10
          # add head and query dimension.
          attention_bias = jnp.expand_dims(attention_bias, -2)
          attention_bias = jnp.expand_dims(attention_bias, -2)
      else:
          attention_bias = None
  
      # apply attention
      x = attention_layers.dot_product_attention(
          query,
          key,
          value,
          dtype=self.dtype,
          bias=attention_bias,
          dropout_rate=self.dropout_rate,
          broadcast_dropout=self.broadcast_dropout,
          dropout_rng=self.make_rng("dropout") if train else None,
          deterministic=not train,
          capture_attention_weights=True,
      )
  
      # 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,
          name="out",
      )(x)
  
      return out
  
imports: import functools
  from typing import Any, Callable, Dict, Tuple, Optional
  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.projects.baselines import resnet
  from scenic.projects.baselines.detr import detr_base_model
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) and returns: out, unknown