code: def init_vit_from_train_state(
      train_state: Any,
      restored_train_state: Any,
      model_cfg: ml_collections.ConfigDict,
      restored_model_cfg: ml_collections.ConfigDict,
  ) -> Any:
      """Updates the train_state with data from restored_train_state.
  
      This function is written to be used for 'fine-tuning' experiments. Here, we
      do some surgery to support larger resolutions (longer sequence length) in
      the transformer block, with respect to the learned pos-embeddings.
  
      The function supports train_states using either Optax or flax.optim (which
      has been deprecated, and will be removed from Scenic.)
  
      Args:
        train_state: A raw TrainState for the model.
        restored_train_state: A TrainState that is loaded with parameters/state of a
          pretrained model.
        model_cfg: Configuration of the model. Usually used for some asserts.
        restored_model_cfg: Configuration of the model from which the
          restored_train_state come from. Usually used for some asserts.
  
      Returns:
        Updated train_state.
      """
      if hasattr(train_state, "optimizer"):
          # TODO(dehghani): Remove support for flax optim.
          params = flax.core.unfreeze(train_state.optimizer.target)
          restored_params = flax.core.unfreeze(
              restored_train_state.optimizer.target
          )
          _merge_params(
              params, restored_params, model_cfg, restored_model_cfg
          )
          return train_state.replace(
              optimizer=train_state.optimizer.replace(
                  target=flax.core.freeze(params)
              )
          )
      else:
          params = flax.core.unfreeze(train_state.params)
          restored_params = flax.core.unfreeze(
              restored_train_state.params
          )
          _merge_params(
              params, restored_params, model_cfg, restored_model_cfg
          )
          return train_state.replace(params=flax.core.freeze(params))
  
imports: from typing import Any, Callable, Optional, Sequence
  from absl import logging
  import flax
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.model_lib.base_models.multilabel_classification_model import MultiLabelClassificationModel
  from scenic.model_lib.layers import attention_layers
  from scenic.model_lib.layers import nn_layers
  import scipy
  from tensorflow.io import gfile
label_desc: Use image node_id 3 for calling a global function with example usage: init_vit_from_train_state(train_state, restored_train_state, model_cfg, restored_model_cfg) and returns: train_state, train_state