class_docstr: DETR Transformer decoder block.
  
  Attributes:
    num_heads: Number of heads.
    qkv_dim: Dimension of the query/key/value.
    mlp_dim: Dimension of the mlp on top of attention block.
    pre_norm: If use LayerNorm before attention/mlp blocks.
    dropout_rate:Dropout rate.
    attention_dropout_rate:Dropout rate for attention weights.
    dtype: Data type of the computation (default: float32).
class_docstr_tok: ["DETR","Transformer","decoder","block",".","Attributes",":","num_heads",":","Number","of","heads",".","qkv_dim",":","Dimension","of","the","query\/key\/value",".","mlp_dim",":","Dimension","of","the","mlp","on","top","of","attention","block",".","pre_norm",":","If","use","LayerNorm","before","attention\/mlp","blocks",".","dropout_rate",":","Dropout","rate",".","attention_dropout_rate",":","Dropout","rate","for","attention","weights",".","dtype",":","Data","type","of","the","computation","(","default",":","float32",")","."]
docstr: Applies DecoderBlock module.
  
  Args:
    obj_queries: Input data for decoder.
    encoder_output: Output of encoder, which are encoded inputs.
    pos_embedding: Positional Embedding to be added to the keys in
      cross-attention.
    query_pos_emb: Positional Embedding to be added to the queries.
    key_padding_mask: Binary mask containing 0 for pad tokens in key.
    train: Train or not (to apply dropout)
  
  Returns:
    Output after transformer decoder block.
code: def __call__(
      self,
      obj_queries: jnp.ndarray,
      encoder_output: jnp.ndarray,
      *,
      pos_embedding: Optional[jnp.ndarray] = None,
      query_pos_emb: Optional[jnp.ndarray] = None,
      key_padding_mask: Optional[jnp.ndarray] = None,
      train: bool = False,
  ):
      """Applies DecoderBlock module.
  
      Args:
        obj_queries: Input data for decoder.
        encoder_output: Output of encoder, which are encoded inputs.
        pos_embedding: Positional Embedding to be added to the keys in
          cross-attention.
        query_pos_emb: Positional Embedding to be added to the queries.
        key_padding_mask: Binary mask containing 0 for pad tokens in key.
        train: Train or not (to apply dropout)
  
      Returns:
        Output after transformer decoder block.
      """
  
      assert query_pos_emb is not None, (
          "Given that object_queries are zeros "
          "and not learnable, we should add "
          "learnable query_pos_emb to them."
      )
      # Seems in DETR the self-attention in the first layer basically does
      # nothing, as the  value vector is a zero vector and we add no learnable
      # positional embedding to it!
      self_attn = MultiHeadDotProductAttention(
          num_heads=self.num_heads,
          qkv_features=self.qkv_dim,
          broadcast_dropout=False,
          dropout_rate=self.attention_dropout_rate,
          kernel_init=initializers.xavier_uniform(),
          bias_init=initializers.zeros,
          use_bias=True,
          dtype=self.dtype,
      )
  
      cross_attn = MultiHeadDotProductAttention(
          num_heads=self.num_heads,
          qkv_features=self.qkv_dim,
          broadcast_dropout=False,
          dropout_rate=self.attention_dropout_rate,
          kernel_init=initializers.xavier_uniform(),
          bias_init=initializers.zeros,
          use_bias=True,
          dtype=self.dtype,
      )
  
      mlp = MlpBlock(  # pytype: disable=wrong-arg-types  # jnp-type
          mlp_dim=self.mlp_dim,
          activation_fn=nn.relu,
          dtype=self.dtype,
          dropout_rate=self.dropout_rate,
      )
  
      assert obj_queries.ndim == 3
      if self.pre_norm:
          # self attention block
          x = nn.LayerNorm(dtype=self.dtype)(obj_queries)
          x = self_attn(
              inputs_q=x,
              pos_emb_q=query_pos_emb,
              pos_emb_k=query_pos_emb,
              pos_emb_v=None,
              train=train,
          )
          x = nn.Dropout(rate=self.dropout_rate)(
              x, deterministic=not train
          )
          x = x + obj_queries
          # cross attention block
          y = nn.LayerNorm(dtype=self.dtype)(x)
          y = cross_attn(
              inputs_q=y,
              inputs_kv=encoder_output,
              pos_emb_q=query_pos_emb,
              pos_emb_k=pos_embedding,
              pos_emb_v=None,
              key_padding_mask=key_padding_mask,
              train=train,
          )
          y = nn.Dropout(rate=self.dropout_rate)(
              y, deterministic=not train
          )
          y = y + x
          # mlp block
          z = nn.LayerNorm(dtype=self.dtype)(y)
          z = mlp(z, deterministic=not train)
          out = y + z
  
      else:
          # self attention block
          x = self_attn(
              inputs_q=obj_queries,
              pos_emb_q=query_pos_emb,
              pos_emb_k=query_pos_emb,
              pos_emb_v=None,
              train=train,
          )
          x = nn.Dropout(rate=self.dropout_rate)(
              x, deterministic=not train
          )
          x = x + obj_queries
          x = nn.LayerNorm(dtype=self.dtype)(x)
          # cross attention block
          y = cross_attn(
              inputs_q=x,
              inputs_kv=encoder_output,
              pos_emb_q=query_pos_emb,
              pos_emb_k=pos_embedding,
              pos_emb_v=None,
              key_padding_mask=key_padding_mask,
              train=train,
          )
          y = nn.Dropout(rate=self.dropout_rate)(
              y, deterministic=not train
          )
          y = y + x
          y = nn.LayerNorm(dtype=self.dtype)(y)
          # mlp block
          z = mlp(y, deterministic=not train)
          z = y + z
          out = nn.LayerNorm(dtype=self.dtype)(z)
  
      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 DecoderBlock obj's underlying member method code with example usage: obj.__call__(obj_queries, encoder_output) and returns: out