class_docstr: Basic ViT + T5 vision language model.
class_docstr_tok: ["Basic","ViT","+","T5","vision","language","model","."]
code: def setup(self):
      t5_config = t5_pretrained.CONFIGS[self.config.t5_name]
      self.t5_config = t5_config
      t5_config["dropout_rate"] = self.config.dropout_rate
      self.ndim = t5_config["emb_dim"]
      self.dropout_rate = t5_config["dropout_rate"]
      self.key_dim = self.config.key_dim
      self.dtype = t5_config["dtype"]
      # Shared token embedding for T5 encoder & Decoder
      self.shared_token_embedder = t5_model.t5_layers.Embed(
          num_embeddings=t5_config["vocab_size"],
          features=self.ndim,
          dtype=self.dtype,
          attend_dtype=self.dtype,  # For logit training stability.
          embedding_init=nn.initializers.normal(stddev=1.0),
          one_hot=True,
          name="shared_token_embedder",
      )
      # Pre-Trained Lower T5 Decoder
      self.out_decoder = t5_model.T5Decoder(
          **t5_config,
          shared_embedding=self.shared_token_embedder,
          name="out_decoder",
      )
      # Uni-Modal Text Encoding (Pre-Trained Lower T5 Encoder)
      self.text_encoder = layers.LowerT5Encoder(
          **t5_config,
          num_fusion_layers=self.config.num_fusion_layers,
          shared_embedding=self.shared_token_embedder,
          name="text_encoder",
      )
      # Multi-Modal Fusion Encoder (Pre-Trained Upper T5 Encoder)
      self.fusion_encoder = layers.FusedT5Encoder(
          **t5_config,
          num_fusion_layers=self.config.num_fusion_layers,
          name="fusion_encoder",
      )
      # Visual Encoding (Pre-Trained ViT)
      self.img_encoder = vit_model.Model(
          num_classes=self.ndim,
          dropout=self.dropout_rate,
          name="img_encoder",
          variant=self.config.vit_name,
          head_zeroinit=False,
          dtype=jnp.bfloat16,
          num_frozen_layers=self.config.get(
              "vit_num_frozen_layers", -1
          ),
          pool_type="gap",
      )
      self.dropout = nn.Dropout(rate=0.2)
  
imports: from typing import Any, Dict, Optional
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  from scenic.model_lib.base_models import base_model
  from scenic.projects.knowledge_visual_language.models import constants
  from scenic.projects.knowledge_visual_language.models import layers
  from scenic.projects.knowledge_visual_language.models import losses
  from scenic.projects.knowledge_visual_language.models import metrics
  from scenic.projects.knowledge_visual_language.models import vit
  from scenic.projects.t5 import layers
  from scenic.projects.t5 import model
label_desc: Use image node_id 1 for calling the VisionLanguageModule obj's underlying member method code with example usage: obj.setup() without return types