class_docstr: Layout denoise model.
  
  Attributes:
    modal_ranges: Modal ranges.
    num_classes: Number of classes.
    num_classes: Number of object classes.
    vocab_size: Vocabulary size.
    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 LayoutDETR 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: ["Layout","denoise","model",".","Attributes",":","modal_ranges",":","Modal","ranges",".","num_classes",":","Number","of","classes",".","num_classes",":","Number","of","object","classes",".","vocab_size",":","Vocabulary","size",".","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","LayoutDETR","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 LayoutDETR model on the input.
  
  Args:
    image: Image data.
    obj_mask: A binary mask where valid is 1 and padding is 0.
    desc_id: description token ids.
    resource_id: resource-id token ids.
    name_id: android class name token ids.
    boxes: Object boxes.
    train:  Whether it is training.
    task: task name.
    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'.
docstr_tok: ["Applies","LayoutDETR","model","on","the","input",".","Args",":","image",":","Image","data",".","obj_mask",":","A","binary","mask","where","valid","is","1","and","padding","is","0.","desc_id",":","description","token","ids",".","resource_id",":","resource-id","token","ids",".","name_id",":","android","class","name","token","ids",".","boxes",":","Object","boxes",".","train",":","Whether","it","is","training",".","task",":","task","name",".","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,
      image: jnp.ndarray,
      obj_mask: jnp.ndarray,
      desc_id: jnp.ndarray,
      resource_id: jnp.ndarray,
      name_id: jnp.ndarray,
      boxes: jnp.ndarray,
      train: bool,
      *,
      task: str,
      padding_mask: Optional[jnp.ndarray] = None,
      update_batch_stats: bool = False,
      debug: bool = False,
  ) -> Dict[str, Any]:
      """Applies LayoutDETR model on the input.
  
      Args:
        image: Image data.
        obj_mask: A binary mask where valid is 1 and padding is 0.
        desc_id: description token ids.
        resource_id: resource-id token ids.
        name_id: android class name token ids.
        boxes: Object boxes.
        train:  Whether it is training.
        task: task name.
        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'.
      """
      assert self.hidden_dim == self.transformer_qkv_dim
  
      num_queries = self.modal_ranges[1] - self.modal_ranges[0]
      token_embder = embedding.TokenEmbedding(
          hidden_dim=self.hidden_dim, vocab_size=self.vocab_size
      )
  
      structure_dict = embedding.StructureEmbedding(
          hidden_dim=self.hidden_dim,
          num_queries=num_queries,
          dropout_rate=self.dropout_rate,
      )(
          obj_mask,
          desc_id,
          resource_id,
          name_id,
          boxes,
          # clickable,
          task=task,
          token_embder=token_embder,
          pos_pattern=self.pos_pattern,
          train=train,
      )
  
      if self.model_type == "full":
          # Full model using both image and view hierarchy structure.
          cnn = resnet.ResNet(
              num_outputs=None,
              num_filters=self.backbone_num_filters,
              num_layers=self.backbone_num_layers,
              dtype=self.dtype,
              name="backbone",
          )
          image_dict = embedding.ImageEmbedding(
              hidden_dim=self.hidden_dim,
              backbone_num_filters=self.backbone_num_filters,
              backbone_num_layers=self.backbone_num_layers,
              name="image_embedding",
          )(
              cnn=cnn,
              images=image,
              train=train,
              padding_mask=padding_mask,
              update_batch_stats=update_batch_stats,
          )
          layout_detr = DeTRTransformer(
              num_queries=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,
              normalize_before=self.transformer_normalize_before,
              dropout_rate=self.dropout_rate,
              attention_dropout_rate=self.attention_dropout_rate,
              return_intermediate_dec=self.aux_loss,
          )
          decoder_output = layout_detr(
              image_dict["content_emb"],
              structure_dict["content_emb"],
              padding_mask=image_dict["mask"],
              pos_embedding=image_dict["pos_emb"],
              query_pos_emb=structure_dict["pos_emb"],
              train=train,
          )
      elif self.model_type == "vh_only":
          # Model only using view hierarchy structure.
          layout_vh_only = VHOnlyModel(
              num_heads=self.transformer_num_heads,
              num_encoder_layers=self.transformer_num_encoder_layers,
              qkv_dim=self.transformer_qkv_dim,
              mlp_dim=self.transformer_mlp_dim,
              normalize_before=self.transformer_normalize_before,
              dropout_rate=self.dropout_rate,
              attention_dropout_rate=self.attention_dropout_rate,
          )
          decoder_output = layout_vh_only(
              vh_inputs=structure_dict["content_emb"],
              pos_embedding=structure_dict["pos_emb"],
              padding_mask=structure_dict["mask"],
              train=train,
          )
      elif self.model_type == "mlp":
          # Model only using individual nodes in view hiearchy without structure.
          layout_mlp = MLPModel(
              num_encoder_layers=self.transformer_num_encoder_layers,
              hidden_dim=self.hidden_dim,
              dropout_rate=self.dropout_rate,
          )
          decoder_output = layout_mlp(
              vh_inputs=structure_dict["content_emb"],
              pos_embedding=structure_dict["pos_emb"],
              padding_mask=structure_dict["mask"],
              train=train,
          )
  
      def output_projection(params):
          model_output = params["multimodal_outputs"]
          return output_projection_pamp(model_output, train)
  
      def output_projection_pamp(model_output, train):
          pred_logits = predictor.ObjectClassPredictor(
              num_classes=self.num_classes,
              dropout_rate=self.class_dropout_rate,
          )(model_output, deterministic=not train)
          return {
              "pred_logits": pred_logits,
          }
  
      params = {
          "multimodal_outputs": decoder_output,
          "train": train,
      }
  
      return common.create_output(
          output_projection,
          params=params,
          aux_loss=self.aux_loss,
          layout_model_pamp=output_projection_pamp,
      )
  
imports: from typing import Any, Dict, List, Tuple, Optional
  import flax.linen
  import jax.numpy
  import ml_collections
  from scenic.projects.baselines import resnet
  from scenic.projects.layout_denoise import base_model
  from scenic.projects.layout_denoise.layers import common
  from scenic.projects.layout_denoise.layers import embedding
  from scenic.projects.layout_denoise.layers import predictor
  from scenic.projects.layout_denoise.layers import transformer
label_desc: Use image node_id 1 for calling the LayoutDenoiseModel obj's underlying member method code with example usage: obj.__call__(image, obj_mask, desc_id, resource_id, name_id, boxes, train) and returns: common, output_projection_pamp, dict