class_docstr: AdaTape Transformer model for multi-label classification task.
class_docstr_tok: ["AdaTape","Transformer","model","for","multi-label","classification","task","."]
code: def build_flax_model(self) -> nn.Module:
      model_dtype = getattr(
          jnp, self.config.get("model_dtype_str", "float32")
      )
      # For the parity task, we use AdaTapeParity as our model.
      if self.config.model_name == "adatape-parity":
          return AdaTapeParity(
              num_classes=self.dataset_meta_data["num_classes"],
              mlp_dim=self.config.model.mlp_dim,
              num_layers=self.config.model.num_layers,
              num_heads=self.config.model.num_heads,
              patches=self.config.model.patches,
              ac_config=self.config.model.ac_config,
              hidden_size=self.config.model.hidden_size,
              classifier=self.config.model.classifier,
              dropout_rate=self.config.model.get("dropout_rate", 0.1),
              attention_dropout_rate=self.config.model.get(
                  "attention_dropout_rate", 0.1
              ),
              stochastic_depth=self.config.model.get(
                  "stochastic_depth", 0.0
              ),
              dtype=model_dtype,
          )
      else:
          return AdaTapeViT(
              num_classes=self.dataset_meta_data["num_classes"],
              mlp_dim=self.config.model.mlp_dim,
              num_layers=self.config.model.num_layers,
              num_heads=self.config.model.num_heads,
              patches=self.config.model.patches,
              ac_config=self.config.model.ac_config,
              hidden_size=self.config.model.hidden_size,
              classifier=self.config.model.classifier,
              dropout_rate=self.config.model.get("dropout_rate", 0.1),
              attention_dropout_rate=self.config.model.get(
                  "attention_dropout_rate", 0.1
              ),
              stochastic_depth=self.config.model.get(
                  "stochastic_depth", 0.0
              ),
              dtype=model_dtype,
          )
  
imports: from typing import Any, Optional
  from absl import logging
  import flax
  import flax.linen
  from flax.training import common_utils
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.common_lib import debug_utils
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import model_utils
  from scenic.model_lib.base_models.multilabel_classification_model import MultiLabelClassificationModel
  from scenic.model_lib.layers import nn_layers
  from scenic.projects.adatape import layers
  from scenic.projects.baselines import vit
label_desc: Use image node_id 1 for calling the AdaTapeMultiLabelClassificationModel obj's underlying member method code with example usage: obj.build_flax_model() and returns: AdaTapeParity, AdaTapeViT