class_docstr: DeformableDETR.
  
  Attributes:
    num_classes: Number of classes to predict.
    embed_dim: Size of the hidden embedding dimension.
    embed_dim: Size of the hidden embedding dimension for encoder.
    num_heads: Number of heads.
    num_queries: Number of object queries.
    num_enc_layers: Number of encoder layers.
    num_dec_layers: Number of decoder layers.
    num_feature_levels: Number of feature levels/scales.
    num_enc_points: Number of encoder points in deformable attention.
    num_dec_points: Number of decoder points in deformable attention.
    transformer_ffn_dim: Transformers feed-forward/MLP dimension.
    backbone_num_filters: Number of filters for Resnet.
    backbone_num_layers: Number of layers for Resnet.
    dropout: Dropout rate.
    compiler_config: Compiler configuration.
    dtype: Data type of the computation (default: float32).
docstr: Perform multi-scale DeformableDETR.
  
  Args:
    inputs: [bs, h, w, c] input data.
    train: Whether it is training.
    padding_mask: [bs, h, w] of bools 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: Necessary for scenic api.
  
  Returns:
    Output:
      'pred_logits' - [bs, num_queries, num_classes] logits.
      'pred_boxes' - [bs, num_queries, 4] boxes in cxcywh format.
docstr_tok: ["Perform","multi-scale","DeformableDETR",".","Args",":","inputs",":","[","bs",",","h",",","w",",","c","]","input","data",".","train",":","Whether","it","is","training",".","padding_mask",":","[","bs",",","h",",","w","]","of","bools","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",":","Necessary","for","scenic","api",".","Returns",":","Output",":","'pred_logits","'","-","[","bs",",","num_queries",",","num_classes","]","logits",".","'pred_boxes","'","-","[","bs",",","num_queries",",","4","]","boxes","in","cxcywh","format","."]
code: def __call__(
      self,
      inputs: jnp.ndarray,
      train: bool = False,
      *,
      padding_mask: Optional[jnp.ndarray] = None,
      update_batch_stats: bool = False,
      debug: bool = False,
  ) -> Dict[str, jnp.ndarray]:
      """Perform multi-scale DeformableDETR.
  
      Args:
        inputs: [bs, h, w, c] input data.
        train: Whether it is training.
        padding_mask: [bs, h, w] of bools 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: Necessary for scenic api.
  
      Returns:
        Output:
          'pred_logits' - [bs, num_queries, num_classes] logits.
          'pred_boxes' - [bs, num_queries, 4] boxes in cxcywh format.
      """
      del debug
      if padding_mask is not None:
          padding_mask = padding_mask.astype(bool)
  
      features, pad_masks, pos_embeds = DeformableDETRBackbone(
          embed_dim=self.enc_embed_dim,
          num_layers=self.backbone_num_layers,
          num_filters=self.backbone_num_filters,
          num_feature_levels=min(self.num_feature_levels, 3),
          dtype=self.dtype,
          name="backbone",
      )(inputs, train=update_batch_stats, padding_mask=padding_mask)
  
      projs = [
          InputProj(
              embed_dim=self.enc_embed_dim, name=f"input_proj{idx}"
          )(x)
          for idx, x in enumerate(features)
      ]
  
      # Add any additional feature scales beyond the last features from backbone.
      x = features[-1]
      for i in range(len(features), self.num_feature_levels):
          x = InputProj(
              embed_dim=self.enc_embed_dim,
              kernel_size=3,
              stride=2,
              padding=1,
              name=f"input_proj{i}",
          )(x)
          projs.append(x)
          pad_masks.append(
              mask_for_shape(x.shape, pad_mask=padding_mask)
          )
          pos_embeds.append(
              InputPosEmbeddingSine(hidden_dim=self.enc_embed_dim)(
                  pad_masks[-1]
              )
          )
  
      # Create shared bbox predictors.
      bbox_embeds = []
      for layer_idx in range(self.num_dec_layers):
          bbox = BBoxCoordPredictor(
              mlp_dim=self.embed_dim,
              num_layers=3,
              use_sigmoid=False,
              name=f"bbox_embed{layer_idx}",
          )
          bbox_embeds.append(bbox)
  
      x, ref_points, dec_init_ref_points = DeformableDETRTransformer(
          enc_embed_dim=self.enc_embed_dim,
          embed_dim=self.embed_dim,
          num_heads=self.num_heads,
          num_queries=self.num_queries,
          num_enc_layers=self.num_enc_layers,
          num_dec_layers=self.num_dec_layers,
          ffn_dim=self.transformer_ffn_dim,
          num_enc_points=self.num_enc_points,
          num_dec_points=self.num_dec_points,
          bbox_embeds=bbox_embeds,
          name="transformer",
          dropout=self.dropout,
          compiler_config=self.compiler_config,
          dtype=self.dtype,
      )(
          inputs=projs,
          pad_masks=pad_masks,
          pos_embeds=pos_embeds,
          train=train,
      )
  
      assert len(x) == self.num_dec_layers
  
      # Classes and Box coordinates prediction heads.
      pred_logits_by_layer = []
      pred_boxes_by_layer = []
      for layer_idx in range(self.num_dec_layers):
          # Logits.
          logits = ObjectClassPredictor(
              num_classes=self.num_classes,
              name=f"class_embed{layer_idx}",
          )(x[layer_idx])
          pred_logits_by_layer.append(logits)
  
          # Predict box coordinates using intermediate decoder outputs.
          if layer_idx == 0:
              level_ref = dec_init_ref_points
          else:
              level_ref = ref_points[layer_idx - 1]
          level_ref = inverse_sigmoid(level_ref)
  
          bbox = bbox_embeds[layer_idx](x[layer_idx])
          if level_ref.shape[-1] == 4:
              bbox += level_ref
          else:
              xy = bbox[..., :2] + level_ref
              bbox = jnp.concatenate([xy, bbox[..., 2:]], -1)
          bbox = nn.sigmoid(bbox)
          pred_boxes_by_layer.append(bbox)
  
      *prev_layers_out, out = jax.tree_util.tree_map(
          lambda logits, boxes: dict(
              pred_logits=logits, pred_boxes=boxes
          ),
          pred_logits_by_layer,
          pred_boxes_by_layer,
      )
      out["aux_outputs"] = prev_layers_out
      return out
  
imports: import math
  from typing import Any, Dict, Optional, Sequence, Tuple, Union
  from absl import logging
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  from scenic.model_lib import matchers
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import box_utils
  from scenic.model_lib.base_models import model_utils
  from scenic.projects.baselines.deformable_detr.backbone import DeformableDETRBackbone
  from scenic.projects.baselines.deformable_detr.backbone import InputPosEmbeddingSine
  from scenic.projects.baselines.deformable_detr.backbone import mask_for_shape
  from scenic.projects.baselines.deformable_detr.deformable_transformer import BBoxCoordPredictor
  from scenic.projects.baselines.deformable_detr.deformable_transformer import DeformableDETRTransformer
  from scenic.projects.baselines.deformable_detr.deformable_transformer import inverse_sigmoid
  from scenic.projects.baselines.deformable_detr.deformable_transformer import pytorch_kernel_init
label_desc: Use image node_id 1 for calling the DeformableDETR obj's underlying member method code with example usage: obj.__call__(inputs, train) and returns: out