code: def initialize_from_polyvit_train_state(
      train_state: Any,
      restored_train_state: Any,
      tokenizer_to_init_from: Optional[str] = None,
      tokenizer_to_init: Optional[str] = None,
      resolution_to_init: Optional[Any] = None,
      initialize_heads: bool = False,
  ) -> Any:
      """Initializes PolyViT with other PolyViT body."""
  
      params = flax.core.unfreeze(train_state.params)
      restored_params = restored_train_state.params
      restored_params = flax.core.unfreeze(restored_params)
  
      for m_key, m_params in restored_params.items():
          if m_key == "vit_encoder":
              params[m_key] = m_params
          elif m_key == "tokenizer":
              # Tokenizer(s) initialization.
              if tokenizer_to_init_from is None:
                  # Only adding tokenizers which are also in `params`.
                  for tm_key, tm_params in m_params.items():
                      if tm_key in params[m_key]:
                          params[m_key][tm_key] = tm_params
              else:
                  # Initializing `tokenizer_to_init` from `tokenizer_to_init_from` and
                  # adapting all needed shapes.
                  params["tokenizer"][tokenizer_to_init][
                      "cls"
                  ] = m_params[tokenizer_to_init_from]["cls"]
                  params["tokenizer"][tokenizer_to_init]["embedding"][
                      "bias"
                  ] = m_params[tokenizer_to_init_from]["embedding"][
                      "bias"
                  ]
                  # Below we assume that that patch h x w is the same for
                  # tokenizer_to_init_from and for tokenizer_to_init.
                  if tokenizer_to_init_from == "tokenizer3d":
                      params["tokenizer"][tokenizer_to_init][
                          "embedding"
                      ]["kernel"] = m_params["tokenizer3d"][
                          "embedding"
                      ][
                          "kernel"
                      ].sum(
                          axis=0
                      )
                  elif tokenizer_to_init == "tokenizer3d":
                      # This is hardcoded since only this shape was used in linear eval of
                      # video tasks.
                      kernel3d = np.zeros([4, 16, 16, 3, 768])
                      # Initializing the middle frame embedding as in ViViT.
                      kernel3d[2, :, :] = m_params[
                          tokenizer_to_init_from
                      ]["embedding"]["kernel"]
                      params["tokenizer"]["tokenizer3d"]["embedding"][
                          "kernel"
                      ] = kernel3d
                  else:
                      params["tokenizer"][tokenizer_to_init][
                          "embedding"
                      ]["kernel"] = m_params[tokenizer_to_init_from][
                          "embedding"
                      ][
                          "kernel"
                      ]
  
                  # Positional embedding initialization.
                  pos_embedding = m_params[tokenizer_to_init_from][
                      "posembed_input"
                  ]["pos_embedding"]
                  # Excluding cls token embedding.
                  new_pos_embedding = pos_embedding[:, 1:]
                  if tokenizer_to_init_from == "tokenizer3d":
                      # This is hardcoded, 8 = 32 / 4 (number of frames in the input /
                      # number of frames in the patch).
                      new_pos_embedding = new_pos_embedding.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 image resolution is
                      # 384 x 384, video frame resolution is 224 x 224 and audio spectrogram
                      # resolution is 800 x 128. We assume that patch size is 16 x 16.
                      if tokenizer_to_init_from == "tokenizer2d":
                          new_pos_embedding = new_pos_embedding.reshape(
                              24, 24, -1
                          )
                          zoom = (
                              resolution_to_init[0] / 384,
                              resolution_to_init[1] / 384,
                              1,
                          )
                      elif tokenizer_to_init_from == "tokenizer_spec":
                          new_pos_embedding = new_pos_embedding.reshape(
                              50, 8, -1
                          )
                          zoom = (
                              resolution_to_init[0] / 800,
                              resolution_to_init[1] / 128,
                              1,
                          )
                      else:
                          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,
                      )
                  if tokenizer_to_init == "tokenizer3d":
                      # This is hardcoded, 8 = 32 / 4 (number of frames in the input /
                      # number of frames in the patch).
                      new_pos_embedding = np.tile(
                          new_pos_embedding, (1, 8, 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 initialize_heads and m_key in params:
              # Initializing heads if needed.
              params[m_key] = 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 1 for calling a global function with example usage: initialize_from_polyvit_train_state(train_state, restored_train_state, tokenizer_to_init_from, tokenizer_to_init, resolution_to_init, initialize_heads) and returns: train_state