class_docstr: Plain Vision Transformer.
class_docstr_tok: ["Plain","Vision","Transformer","."]
code: def build_flax_model(self) -> nn.Module:
      dtype_str = self.config.get("model_dtype_str", "float32")
      if dtype_str != "float32":
          raise ValueError(
              "`dtype` argument is not propagated properly "
              "in the current implementation, so only "
              "`float32` is supported for now."
          )
      return ViT(
          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,
          positional_embedding=self.config.model.positional_embedding,
          representation_size=self.config.model.representation_size,
          patches=self.config.model.patches,
          hidden_size=self.config.model.hidden_size,
          classifier=self.config.model.classifier,
          dropout_rate=self.config.model.dropout_rate,
          dtype=getattr(jnp, dtype_str),
      )
  
imports: import functools
  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.model_lib.base_models import base_model
  from scenic.model_lib.base_models import classification_model
  from scenic.model_lib.base_models import model_utils
  from scenic.model_lib.base_models import multilabel_classification_model
  from scenic.model_lib.layers import nn_layers
  import scipy
label_desc: Use image node_id 1 for calling the PlainViT obj's underlying member method code with example usage: obj.build_flax_model() and returns: ViT