code: def init_posemb(
      to_params,
      from_params,
      config,
      restored_model_cfg,
      positional_embed_size_change,
  ):
      """Initialize the positional embeddings."""
      with_cls_token, num_video_frames = get_cls_token_and_video_frames(
          config
      )
      restored_with_cls_token, _ = get_cls_token_and_video_frames(
          restored_model_cfg
      )
      if config.init_from.restore_positional_embedding:
          posemb = to_params["posembed_input"]["pos_embedding"]
          restored_posemb = from_params["posembed_input"][
              "pos_embedding"
          ]
          if restored_posemb.shape != posemb.shape:
              # Rescale the grid of pos, embeddings.
              # Default parameter shape is (1, N, 768)
              logging.info(
                  "Adapting positional embeddings from %s to %s",
                  restored_posemb.shape,
                  posemb.shape,
              )
              ntok = posemb.shape[1]
              if restored_with_cls_token:
                  # The first token is the CLS token.
                  cls_tok = restored_posemb[:, :1]
                  restored_posemb_grid = restored_posemb[0, 1:]
              else:
                  cls_tok = restored_posemb[:, :0]
                  restored_posemb_grid = restored_posemb[0]
              if with_cls_token:
                  ntok -= 1
              restored_gs = int(np.sqrt(len(restored_posemb_grid)))
              gs = int(np.sqrt(ntok))
              if with_cls_token != restored_with_cls_token:
                  logging.warning(
                      "Only one of target and restored model uses"
                      "classification token"
                  )
                  if restored_gs == gs:
                      # In case the following `if` is not going to run, lets add batch dim:
                      restored_posemb = restored_posemb_grid[None, ...]
  
              if restored_gs != gs:  # We need resolution change.
                  if positional_embed_size_change == "resize":
                      restored_posemb_grid = (
                          vivit_utils.interpolate_positional_embeddings(
                              restored_posemb_grid, ntok
                          )
                      )
  
                  elif positional_embed_size_change == "tile":
                      restored_posemb_grid = (
                          vivit_utils.tile_positional_embeddings(
                              restored_posemb_grid, ntok
                          )
                      )
  
                  elif positional_embed_size_change == "resize_tile":
                      n_frames = (
                          num_video_frames
                          // config.model.modalities.video.patches.size[
                              2
                          ]
                      )
                      tokens_per_frame = ntok // n_frames
                      restored_posemb_grid = (
                          vivit_utils.interpolate_positional_embeddings(
                              restored_posemb_grid, tokens_per_frame
                          )
                      )
                      restored_posemb_grid = restored_posemb_grid[0]
                      restored_posemb_grid = (
                          vivit_utils.tile_positional_embeddings(
                              restored_posemb_grid, ntok
                          )
                      )
  
                  else:
                      raise AssertionError(
                          "Unknown positional embedding size changing method"
                      )
                  # Attach the CLS token again.
                  if with_cls_token:
                      restored_posemb = jnp.array(
                          np.concatenate(
                              [cls_tok, restored_posemb_grid], axis=1
                          )
                      )
                  else:
                      restored_posemb = restored_posemb_grid
  
          to_params["posembed_input"]["pos_embedding"] = restored_posemb
      else:
          logging.info(
              "Not restoring positional encodings from pretrained model"
          )
  
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 5 for calling a global function with example usage: init_posemb(to_params, from_params, config, restored_model_cfg, positional_embed_size_change) without return types