class_docstr: Detection Transformer (DETR) model.
  
  Attributes:
    num_classes: Number of object classes.
    hidden_dim: Hidden dimension of the inputs to the model.
    num_queries: Number of object queries, ie detection slot. This is the
      maximal number of objects DETR can detect in a single image. For COCO,
      DETR paper recommends 100 queries.
    query_emb_size: Size of the embedding learned for object queries.
    transformer_num_heads: Number of transformer heads.
    transformer_num_encoder_layers: Number of transformer encoder layers.
    transformer_num_decoder_layers: Number of transformer decoder layers.
    transformer_qkv_dim: Dimension of the transformer query/key/value.
    transformer_mlp_dim: Dimension of the mlp on top of attention block.
    transformer_normalize_before: If use LayerNorm before attention/mlp blocks.
    backbone_num_filters: Num filters in the ResNet backbone.
    backbone_num_layers: Num layers in the ResNet backbone.
    aux_loss: If train with auxiliary loss.
    dropout_rate:Dropout rate.
    attention_dropout_rate:Attention dropout rate.
    dtype: Data type of the computation (default: float32).
class_docstr_tok: ["Detection","Transformer","(","DETR",")","model",".","Attributes",":","num_classes",":","Number","of","object","classes",".","hidden_dim",":","Hidden","dimension","of","the","inputs","to","the","model",".","num_queries",":","Number","of","object","queries",",","ie","detection","slot",".","This","is","the","maximal","number","of","objects","DETR","can","detect","in","a","single","image",".","For","COCO",",","DETR","paper","recommends","100","queries",".","query_emb_size",":","Size","of","the","embedding","learned","for","object","queries",".","transformer_num_heads",":","Number","of","transformer","heads",".","transformer_num_encoder_layers",":","Number","of","transformer","encoder","layers",".","transformer_num_decoder_layers",":","Number","of","transformer","decoder","layers",".","transformer_qkv_dim",":","Dimension","of","the","transformer","query\/key\/value",".","transformer_mlp_dim",":","Dimension","of","the","mlp","on","top","of","attention","block",".","transformer_normalize_before",":","If","use","LayerNorm","before","attention\/mlp","blocks",".","backbone_num_filters",":","Num","filters","in","the","ResNet","backbone",".","backbone_num_layers",":","Num","layers","in","the","ResNet","backbone",".","aux_loss",":","If","train","with","auxiliary","loss",".","dropout_rate",":","Dropout","rate",".","attention_dropout_rate",":","Attention","dropout","rate",".","dtype",":","Data","type","of","the","computation","(","default",":","float32",")","."]
docstr: Applies DETR model on the input.
  
  Args:
    inputs: Input data.
    train:  Whether it is training.
    padding_mask: Binary matrix with 0 at padded image regions.
    update_batch_stats: Whether update the batch statistics for the BatchNorms
      in the backbone. if None, the value of `train` flag will be used, i.e.
      we update the batch stat if we are in the train mode.
    debug: Whether the debug mode is enabled. debug=True enables model
      specific logging/storing some values using jax.host_callback.
  
  Returns:
    Output: dit; that has 'pred_logits' and 'pred_boxes', and potentially
    'aux_outputs'.
code: def __call__(
      self,
      inputs: jnp.ndarray,
      train: bool,
      *,
      padding_mask: Optional[jnp.ndarray] = None,
      update_batch_stats: bool = False,
      debug: bool = False,
  ) -> Dict[str, Any]:
      """Applies DETR model on the input.
  
      Args:
        inputs: Input data.
        train:  Whether it is training.
        padding_mask: Binary matrix with 0 at padded image regions.
        update_batch_stats: Whether update the batch statistics for the BatchNorms
          in the backbone. if None, the value of `train` flag will be used, i.e.
          we update the batch stat if we are in the train mode.
        debug: Whether the debug mode is enabled. debug=True enables model
          specific logging/storing some values using jax.host_callback.
  
      Returns:
        Output: dit; that has 'pred_logits' and 'pred_boxes', and potentially
        'aux_outputs'.
      """
      # for now, we only support this case
      assert self.hidden_dim == self.transformer_qkv_dim
  
      if update_batch_stats is None:
          update_batch_stats = train
  
      backbone_features = resnet.ResNet(
          num_outputs=None,
          num_filters=self.backbone_num_filters,
          num_layers=self.backbone_num_layers,
          dtype=self.dtype,
          name="backbone",
      )(inputs, train=update_batch_stats)
      x = backbone_features["stage_4"]
  
      bs, h, w, _ = x.shape
  
      if padding_mask is None:
          padding_mask_downsampled = jnp.ones(
              (bs, h, w), dtype=jnp.bool_
          )
      else:
          padding_mask_downsampled = jax.image.resize(
              padding_mask.astype(jnp.float32),
              shape=[bs, h, w],
              method="nearest",
          ).astype(jnp.bool_)
      pos_emb = InputPosEmbeddingSine(hidden_dim=self.hidden_dim)(
          padding_mask_downsampled
      )
  
      query_pos_emb = QueryPosEmbedding(
          hidden_dim=self.hidden_dim, num_queries=self.num_queries
      )()
  
      # project and reshape to 3 dimensions and project
      x = nn.Conv(
          features=self.hidden_dim, kernel_size=(1, 1), strides=(1, 1)
      )(x)
      x = x.reshape(bs, h * w, self.hidden_dim)
      transformer_input = x
  
      return_intermediate = self.aux_loss
      transformer = DETRTransformer(
          num_queries=self.num_queries,
          query_emb_size=self.query_emb_size,
          num_heads=self.transformer_num_heads,
          num_encoder_layers=self.transformer_num_encoder_layers,
          num_decoder_layers=self.transformer_num_decoder_layers,
          qkv_dim=self.transformer_qkv_dim,
          mlp_dim=self.transformer_mlp_dim,
          return_intermediate_dec=return_intermediate,
          normalize_before=self.transformer_normalize_before,
          dropout_rate=self.dropout_rate,
          attention_dropout_rate=self.attention_dropout_rate,
      )
      decoder_output, encoder_output = transformer(
          x,
          padding_mask=jnp.reshape(
              padding_mask_downsampled, [bs, h * w]
          ),
          pos_embedding=pos_emb,
          query_pos_emb=query_pos_emb,
          train=train,
      )
  
      def output_projection(model_output):
          # classification head
          pred_logits = ObjectClassPredictor(
              num_classes=self.num_classes
          )(model_output)
          # bounding box detection head
          pred_boxes = BBoxCoordPredictor(mlp_dim=self.hidden_dim)(
              model_output
          )
          return pred_logits, pred_boxes
  
      if not return_intermediate:
          pred_logits, pred_boxes = output_projection(decoder_output)
          return {"pred_logits": pred_logits, "pred_boxes": pred_boxes}
  
      pred_logits, pred_boxes = jax.vmap(output_projection)(
          decoder_output
      )
      output = {
          "pred_logits": pred_logits[-1],
          "pred_boxes": pred_boxes[-1],
          "transformer_input": transformer_input,
          "backbone_features": backbone_features,
          "encoder_output": encoder_output,
          "decoder_output": decoder_output[-1],
          "padding_mask": padding_mask_downsampled,
      }
  
      if self.aux_loss:
          output["aux_outputs"] = []
          for lgts, bxs in zip(pred_logits[:-1], pred_boxes[:-1]):
              output["aux_outputs"].append(
                  {"pred_logits": lgts, "pred_boxes": bxs}
              )
  
      return output
  
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 DETR obj's underlying member method code with example usage: obj.__call__(inputs, train) and returns: output, pred_logits,  pred_boxes, dict