code: def get_config():
      """Returns the ViT experiment configuration for ImageNet.
  
      This file is a copy of config/imagenet_augreg_matvit_config.py.
      """
      version, patch = VARIANT.split("/")
      config = ml_collections.ConfigDict()
      config.experiment_name = "imagenet-regularized_vit"
      # Dataset.
      config.dataset_name = "bit"
      config.data_dtype_str = "float32"
      config.dataset_configs = ml_collections.ConfigDict()
      config.dataset_configs.dataset = "imagenet2012"
      config.dataset_configs.num_classes = NUM_CLASSES
      config.dataset_configs.train_split = "train"
      config.dataset_configs.val_split = "validation"
      config.dataset_configs.pp_train = {
          "L": (
              "decode|resize(384)"
              "|value_range(-1, 1)"
              f'|onehot({NUM_CLASSES}, key="label", key_result="labels")'
              '|keep("image", "labels")'
          ),
          "B": (
              "decode|resize_small(256)|central_crop(224)"
              "|value_range(-1, 1)"
              f'|onehot({NUM_CLASSES}, key="label", key_result="labels")'
              '|keep("image", "labels")'
          ),
      }[version]
      config.dataset_configs.pp_eval = {
          "L": (
              "decode|resize(384)"
              "|value_range(-1, 1)"
              f'|onehot({NUM_CLASSES}, key="label", key_result="labels")'
              '|keep("image", "labels")'
          ),
          "B": (
              "decode|resize_small(256)|central_crop(224)"
              "|value_range(-1, 1)"
              f'|onehot({NUM_CLASSES}, key="label", key_result="labels")'
              '|keep("image", "labels")'
          ),
      }[version]
  
      config.dataset_configs.prefetch_to_device = 2
      # Shuffle_buffer_size is per host, so small-ish is ok.
      config.dataset_configs.shuffle_buffer_size = 250_000
      # Model.
      config.model_name = "vit_multilabel_classification"
      config.model = ml_collections.ConfigDict()
      config.model.hidden_size = {
          "Ti": 192,
          "S": 384,
          "B": 768,
          "L": 1024,
          "H": 1280,
      }[version]
      config.model.patches = ml_collections.ConfigDict()
      config.model.patches.size = [int(patch), int(patch)]
      config.model.num_heads = {
          "Ti": 3,
          "S": 6,
          "B": 12,
          "L": 16,
          "H": 16,
      }[version]
      config.model.mlp_dim = {
          "Ti": 768,
          "S": 1536,
          "B": 3072,
          "L": 4096,
          "H": 5120,
      }[version]
      config.model.num_layers = {
          "Ti": 12,
          "S": 12,
          "B": 12,
          "L": 24,
          "H": 32,
      }[version]
      config.model.representation_size = None
      config.model.classifier = "token"
      config.model.attention_dropout_rate = 0.0
      config.model.dropout_rate = 0.1
      config.model.stochastic_depth = 0.1
      config.model_dtype_str = "float32"
      # Training.
      config.trainer_name = "classification_trainer"
      config.optimizer = "adam"
      config.optimizer_configs = ml_collections.ConfigDict()
      config.optimizer_configs.beta1 = 0.9
      config.optimizer_configs.beta2 = 0.999
      config.optimizer_configs.weight_decay = 0.1
      config.explicit_weight_decay = None  # No explicit weight decay
      config.l2_decay_factor = None
      config.max_grad_norm = 1.0
      config.label_smoothing = None
      config.num_training_epochs = 300
      config.log_eval_steps = 1000
      config.batch_size = 2048
      config.rng_seed = 42
      config.init_head_bias = -6.9  # -log(1000)
  
      # Learning rate.
      steps_per_epoch = _IMAGENET_TRAIN_SIZE // config.batch_size
      total_steps = config.num_training_epochs * steps_per_epoch
      base_lr = 0.001
      config.lr_configs = ml_collections.ConfigDict()
      config.lr_configs.learning_rate_schedule = "compound"
      config.lr_configs.factors = (
          "constant * cosine_decay * linear_warmup"
      )
      config.lr_configs.warmup_steps = 10_000
      config.lr_configs.steps_per_cycle = total_steps
      config.lr_configs.base_learning_rate = base_lr
      # Mixup.
      config.mixup = ml_collections.ConfigDict()
      config.mixup.bind_to = None
      config.mixup.alpha = 0.5
      # Logging.
      config.write_summary = True
      config.xprof = True  # Profile using xprof.
      config.checkpoint = True  # Do checkpointing.
      config.checkpoint_steps = 5000
      config.debug_train = False  # Debug mode during training.
      config.debug_eval = False  # Debug mode during eval.
      config.m = None  # Placeholder for randaug strength.
      config.l = None  # Placeholder for randaug layers.
      return config
  
file_path: scenic/scenic/projects/matvit/classification_eval_main.py
imports: import functools
  from typing import Callable, Dict, List, Optional, Tuple
  from absl import app
  from absl import flags
  from absl import logging
  from flax import jax_utils
  import flax.linen
  import jax
  import jax.numpy
  import jax.profiler
  import ml_collections
  import numpy
  from scenic.projects.matvit import matvit
  from scenic.train_lib import train_utils
label_desc: Use image node_id 1 for calling a global function with example usage: get_config() and returns: config