class_docstr: DETR Transformer.
  
  Attributes:
    num_queries: Number of object queries. query_emb_size; Size of the embedding
      learned for object queries.
    query_emb_size: Size of the embedding learned for object queries.
    num_heads: Number of heads.
    num_encoder_layers: Number of encoder layers.
    num_decoder_layers: Number of decoder layers.
    qkv_dim: Dimension of the query/key/value.
    mlp_dim: Dimension of the mlp on top of attention block.
    return_intermediate_dec: If return the outputs from intermediate layers of
      the decoder.
    normalize_before: 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",".","Attributes",":","num_queries",":","Number","of","object","queries",".","query_emb_size",";","Size","of","the","embedding","learned","for","object","queries",".","query_emb_size",":","Size","of","the","embedding","learned","for","object","queries",".","num_heads",":","Number","of","heads",".","num_encoder_layers",":","Number","of","encoder","layers",".","num_decoder_layers",":","Number","of","decoder","layers",".","qkv_dim",":","Dimension","of","the","query\/key\/value",".","mlp_dim",":","Dimension","of","the","mlp","on","top","of","attention","block",".","return_intermediate_dec",":","If","return","the","outputs","from","intermediate","layers","of","the","decoder",".","normalize_before",":","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 DETRTransformer on the inputs.
  
  Args:
    inputs: Input data.
    padding_mask: Binary mask containing 0 for padding tokens.
    pos_embedding: Positional Embedding to be added to the inputs.
    query_pos_emb: Positional Embedding to be added to the queries.
    train: Whether it is training.
  
  Returns:
    Output of the DETR transformer and output of the encoder.
code: def __call__(
      self,
      inputs: jnp.ndarray,
      *,
      padding_mask: Optional[jnp.ndarray] = None,
      pos_embedding: Optional[jnp.ndarray] = None,
      query_pos_emb: Optional[jnp.ndarray] = None,
      train: bool = False,
  ) -> Tuple[jnp.ndarray, jnp.ndarray]:
      """Applies DETRTransformer on the inputs.
  
      Args:
        inputs: Input data.
        padding_mask: Binary mask containing 0 for padding tokens.
        pos_embedding: Positional Embedding to be added to the inputs.
        query_pos_emb: Positional Embedding to be added to the queries.
        train: Whether it is training.
  
      Returns:
        Output of the DETR transformer and output of the encoder.
      """
      encoder_norm = nn.LayerNorm() if self.normalize_before else None
      encoded = Encoder(
          num_heads=self.num_heads,
          num_layers=self.num_encoder_layers,
          qkv_dim=self.qkv_dim,
          mlp_dim=self.mlp_dim,
          normalize_before=self.normalize_before,
          norm=encoder_norm,
          dropout_rate=self.dropout_rate,
          attention_dropout_rate=self.attention_dropout_rate,
          dtype=self.dtype,
          name="encoder",
      )(
          inputs,
          padding_mask=padding_mask,
          pos_embedding=pos_embedding,
          train=train,
      )
  
      query_dim = self.query_emb_size or inputs.shape[-1]
      obj_query_shape = tuple(
          [inputs.shape[0], self.num_queries, query_dim]
      )
      # Note that we always learn query_pos_embed, so we simply use constant
      # zero vectors for obj_queries and later when applying attention, we have:
      # query = query_pos_embed + obj_queries
      obj_queries = jnp.zeros(obj_query_shape)
  
      decoder_norm = nn.LayerNorm()
      output = Decoder(
          num_heads=self.num_heads,
          num_layers=self.num_decoder_layers,
          qkv_dim=self.qkv_dim,
          mlp_dim=self.mlp_dim,
          normalize_before=self.normalize_before,
          return_intermediate=self.return_intermediate_dec,
          dropout_rate=self.dropout_rate,
          attention_dropout_rate=self.attention_dropout_rate,
          norm=decoder_norm,
          dtype=self.dtype,
          name="decoder",
      )(
          obj_queries,
          encoded,
          key_padding_mask=padding_mask,
          pos_embedding=pos_embedding,
          query_pos_emb=query_pos_emb,
          train=train,
      )
      return output, encoded
  
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 DETRTransformer obj's underlying member method code with example usage: obj.__call__(inputs) and returns: output,  encoded