code: def initialize_from_vivit_train_state(
      train_state: Any,
      restored_train_state: Any,
      tokenizer_to_init: str = "tokenizer3d",
      resolution_to_init: Optional[Any] = None,
      initialize_head: bool = False,
  ) -> Any:
      """Initializes PolyViT with ViViT body."""
  
      params = flax.core.unfreeze(train_state.params)
      restored_params = flax.core.unfreeze(restored_train_state.params)
  
      for m_key, m_params in restored_params.items():
          # Not initializing heads.
          if m_key == "Transformer":
              for tm_key, tm_params in m_params.items():
                  if tm_key == "posembed_input":
                      # Initializing positional embeddings.
                      if tokenizer_to_init == "tokenizer3d":
                          params["tokenizer"]["tokenizer3d"][
                              "posembed_input"
                          ] = tm_params
                      else:
                          # Adapting positional embedding shapes if needed.
                          pos_embedding = tm_params["pos_embedding"]
                          # Excluding cls token embedding and reshaping into per-frame
                          # embeddings. 8 = 32 / 4 (number of frames in the input /
                          # number of frames in the patch).
                          new_pos_embedding = pos_embedding[
                              :, 1:
                          ].reshape(1, 8, -1, pos_embedding.shape[2])
                          # Averaging per-frame positional embeddings to obtain a 2d
                          # initialization.
                          new_pos_embedding = new_pos_embedding.mean(
                              axis=1
                          )
                          # 2d interpolation if needed.
                          if resolution_to_init is not None:
                              # This is hardcoded. We assume that video frame resolution is
                              # 224 x 224 and patch size is 16 x 16.
                              new_pos_embedding = (
                                  new_pos_embedding.reshape(14, 14, -1)
                              )
                              zoom = (
                                  resolution_to_init[0] / 224,
                                  resolution_to_init[1] / 224,
                                  1,
                              )
                              new_pos_embedding = scipy.ndimage.zoom(
                                  new_pos_embedding, zoom, order=1
                              )
                              new_pos_embedding = (
                                  new_pos_embedding.reshape(
                                      1,
                                      new_pos_embedding.shape[0]
                                      * new_pos_embedding.shape[1],
                                      -1,
                                  )
                              )
                          # Concatenating with the cls embedding.
                          new_pos_embedding = np.concatenate(
                              [pos_embedding[:, :1], new_pos_embedding],
                              axis=1,
                          )
                          params["tokenizer"][tokenizer_to_init][
                              "posembed_input"
                          ]["pos_embedding"] = new_pos_embedding
                  elif tm_key in params["vit_encoder"]:
                      params["vit_encoder"][tm_key] = tm_params
          elif m_key == "cls":
              params["tokenizer"][tokenizer_to_init]["cls"] = m_params
          elif m_key == "embedding":
              # Patch embedding initialization.
              if tokenizer_to_init == "tokenizer3d":
                  params["tokenizer"]["tokenizer3d"][
                      "embedding"
                  ] = m_params
              else:
                  params["tokenizer"][tokenizer_to_init]["embedding"][
                      "bias"
                  ] = m_params["bias"]
                  # 2d patch embedding is a sum along the frame dimension of the 3d video
                  # patch embedding.
                  params["tokenizer"][tokenizer_to_init]["embedding"][
                      "kernel"
                  ] = m_params["kernel"].sum(axis=0)
          elif m_key == "output_projection" and initialize_head:
              # Initializing head if needed.
              head_name = [
                  x
                  for x in params.keys()
                  if x not in ["tokenizer", "vit_encoder"]
              ][0]
              params[head_name]["output_projection"] = m_params
  
      return train_state.replace(params=flax.core.freeze(params))
  
file_path: scenic/scenic/projects/polyvit/model_utils.py
imports: from typing import Any, Optional
  from absl import logging
  import flax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.common_lib import debug_utils
  import scenic.projects.mbt.model_utils
  import scenic.projects.vivit.model_utils
  import scipy
label_desc: Use image node_id 3 for calling a global function with example usage: initialize_from_vivit_train_state(train_state, restored_train_state, tokenizer_to_init, resolution_to_init, initialize_head) and returns: train_state