class_docstr: PolyVit model.
class_docstr_tok: ["PolyVit","model","."]
code: def build_flax_model(self):
      model_dtype = getattr(
          jnp, self.config.get("model_dtype_str", "float32")
      )
  
      add_cls_token, _ = model_utils.get_cls_token_and_video_frames(
          self.config
      )
  
      for dataset_name, meta_data in self.dataset_meta_data.items():
          if meta_data["task"] in [
              polyvit_base_model.Task.LABEL,
              polyvit_base_model.Task.MULTILABEL,
              polyvit_base_model.Task.MULTIHEADLABEL,
          ]:
              with self.config.unlocked():
                  self.config.model.heads.label.get(
                      dataset_name
                  ).num_classes = meta_data["num_classes"]
          if meta_data["task"] in [polyvit_base_model.Task.BOW]:
              with self.config.unlocked():
                  self.config.model.heads.bow.get(
                      dataset_name
                  ).vocab_size = meta_data["vocab_size"]
  
      with self.config.unlocked():
          modalities_cfg = self.config.model.modalities
          encoder_cfg = self.config.model.encoder
          modalities_cfg.mlp_dim = encoder_cfg.mlp_dim
          if modalities_cfg.get("num_layers") is None:
              modalities_cfg.num_layers = 0
          modalities_cfg.num_heads = encoder_cfg.num_heads
          modalities_cfg.dropout_rate = encoder_cfg.dropout_rate
          modalities_cfg.attention_dropout_rate = (
              encoder_cfg.attention_dropout_rate
          )
          modalities_cfg.add_cls_token = add_cls_token
          encoder_cfg.num_tokenizer_layers = modalities_cfg.num_layers
  
      stochastic_droplayer_config = self.config.get(
          "stochastic_droplayer_rates"
      )
  
      return PolyVit(
          modalities_config=self.config.model.modalities,
          encoder_config=self.config.model.encoder,
          heads_config=self.config.model.heads,
          stochastic_droplayer_config=stochastic_droplayer_config,
          dtype=model_dtype,
      )
  
file_path: scenic/scenic/projects/polyvit/model.py
imports: from typing import Any, Callable, Optional, Sequence
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  from scenic.projects.polyvit import layers
  from scenic.projects.polyvit import model_utils
  from scenic.projects.polyvit import polyvit_base_model
label_desc: Use image node_id 1 for calling the PolyVitModel obj's underlying member method code with example usage: obj.build_flax_model() and returns: PolyVit