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 writen 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.
  
      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.
      """
      params = flax.core.unfreeze(train_state.optimizer.target)
      restored_params = flax.core.unfreeze(
          restored_train_state.optimizer.target
      )
  
      # Start moving parameters, one-by-one and apply changes if needed.
      for m_key, m_params in restored_params.items():
          if m_key == "output_projection":
              # For the classifier head, we use a the randomly initialized params and
              #   ignore the the one from pretrained model.
              pass
  
          elif m_key == "pre_logits":
              if model_cfg.model.representation_size is None:
                  # We don't have representation_size in the new model, so let's ignore
                  #   it from the pretained model, in case it has it.
                  # Note, removing the key from the dictionary is necessary to prevent
                  #   obscure errors from the Flax optimizer.
                  params.pop(m_key, None)
              else:
                  assert restored_model_cfg.model.representation_size
                  params[m_key] = m_params
  
          elif m_key == "Transformer":
              for tm_key, tm_params in m_params.items():
                  if (
                      tm_key == "posembed_input"
                  ):  # Might need resolution change.
                      posemb = params[m_key]["posembed_input"][
                          "pos_embedding"
                      ]
                      restored_posemb = m_params["posembed_input"][
                          "pos_embedding"
                      ]
  
                      if restored_posemb.shape != posemb.shape:
                          # Rescale the grid of pos, embeddings: param shape is (1, N, d).
                          logging.info(
                              "Resized variant: %s to %s",
                              restored_posemb.shape,
                              posemb.shape,
                          )
                          ntok = posemb.shape[1]
                          if (
                              restored_model_cfg.model.classifier
                              == "token"
                          ):
                              # The first token is the CLS token.
                              cls_tok = restored_posemb[:, :1]
                              restored_posemb_grid = restored_posemb[
                                  0, 1:
                              ]
                              ntok -= 1
                          else:
                              cls_tok = restored_posemb[:, :0]
                              restored_posemb_grid = restored_posemb[0]
  
                          restored_gs = int(
                              np.sqrt(len(restored_posemb_grid))
                          )
                          gs = int(np.sqrt(ntok))
                          if (
                              restored_gs != gs
                          ):  # We need resolution change.
                              logging.info(
                                  "Grid-size from %s to %s.",
                                  restored_gs,
                                  gs,
                              )
                              restored_posemb_grid = (
                                  restored_posemb_grid.reshape(
                                      restored_gs, restored_gs, -1
                                  )
                              )
                              zoom = (
                                  gs / restored_gs,
                                  gs / restored_gs,
                                  1,
                              )
                              restored_posemb_grid = scipy.ndimage.zoom(
                                  restored_posemb_grid, zoom, order=1
                              )
                              restored_posemb_grid = (
                                  restored_posemb_grid.reshape(
                                      1, gs * gs, -1
                                  )
                              )
                              # Attache the CLS token again.
                              restored_posemb = jnp.array(
                                  np.concatenate(
                                      [cls_tok, restored_posemb_grid],
                                      axis=1,
                                  )
                              )
  
                      params[m_key][tm_key][
                          "pos_embedding"
                      ] = restored_posemb
                  else:  # Other parameters of the Transformer encoder.
                      params[m_key][tm_key] = tm_params
  
          else:
              # Use the rest as they are in the pretrianed model.
              params[m_key] = m_params
  
      return train_state.replace(
          optimizer=train_state.optimizer.replace(
              target=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
label_desc: Use image node_id 1 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