code: def init_from_vit_checkpoints(
      config: ml_collections.ConfigDict,
      model: model_lib.MTVClassificationModel,
      train_state: train_utils.TrainState,
  ) -> train_utils.TrainState:
      """Initialize train state from ViT checkpoints."""
      if config.init_from.get("model_cfg") is None:
          logging.info(
              "model_cfg is empty. Use current model's classifier."
          )
          restored_model_cfgs = []
          for _ in range(len(config.model.view_configs)):
              restored_model_cfgs.append(
                  ml_collections.ConfigDict(
                      {"model": {"classifier": config.model.classifier}}
                  )
              )
      else:
          restored_model_cfgs = config.init_from.model_cfg
      init_checkpoint_paths = config.init_from.get("checkpoint_path")
      assert len(init_checkpoint_paths) == len(
          config.model.view_configs
      ), (
          "Number of initial checkpoint paths must match with the number of view "
          "configs."
      )
      checkpoint_formats = config.init_from.get("checkpoint_formats")
      checkpoint_formats = (
          ["scenic"] * len(init_checkpoint_paths)
          if checkpoint_formats is None
          else checkpoint_formats
      )
      assert len(checkpoint_formats) == len(
          init_checkpoint_paths
      ), "The lengths of checkpoint_formats and init_checkpoint_paths must match."
      assert set(checkpoint_formats).issubset(
          {"scenic", "big_vision"}
      ), "Only scenic and big_vision formats are supported."
      restored_train_states = []
      for path, checkpoint_format in zip(
          init_checkpoint_paths, checkpoint_formats
      ):
          if checkpoint_format == "big_vision":
              restored_train_states.append(
                  pretrain_utils.convert_big_vision_to_scenic_checkpoint(
                      path, train_state
                  )
              )
          else:
              restored_train_states.append(
                  pretrain_utils.restore_pretrained_checkpoint(
                      path, train_state, assert_exist=True
                  )
              )
      return model.init_from_vit_train_states(
          train_state,
          restored_train_states,
          restored_model_cfgs,
          checkpoint_formats,
      )
  
imports: import copy
  import functools
  import os
  from typing import Any, Dict, Tuple
  from absl import logging
  from clu import metric_writers
  from clu import periodic_actions
  from clu import platform
  from flax import jax_utils
  import jax
  import jax.numpy
  import jax.profiler
  import ml_collections
  import numpy
  from scenic.dataset_lib import dataset_utils
  from scenic.projects.mtv import model
  from scenic.projects.mtv import train_utils
  from scenic.projects.vivit import evaluation_lib
  from scenic.projects.vivit import train_utils
  from scenic.train_lib import lr_schedules
  from scenic.train_lib import optimizers
  from scenic.train_lib import pretrain_utils
  from scenic.train_lib import train_utils
  import tensorflow
label_desc: Use image node_id 2 for calling a global function with example usage: init_from_vit_checkpoints(config, model, train_state) and returns: model