code: def initialize_from_unloc_parameters(
      params: Dict[str, Any],
      restored_params: Dict[str, Any],
      skip_regex: Optional[str] = None,
  ) -> Dict[str, Any]:
      """Initialize model parameters from an UnLoc model.
  
      Args:
        params: The parameters of the model.
        restored_params: Restored parameters from the given pretrained checkpoint.
        skip_regex: Regular expression of parameters to skip loading.
  
      Returns:
        Initialized parameters of the current model.
      """
  
      restored_flat = flax.traverse_util.flatten_dict(
          dict(restored_params), keep_empty_nodes=True, sep="/"
      )
      model_flat = flax.traverse_util.flatten_dict(
          dict(params), keep_empty_nodes=True, sep="/"
      )
      logging.info("model_flat keys: %s", list(model_flat.keys()))
  
      for m_key, m_params in restored_flat.items():
          if m_key not in model_flat:
              logging.warning(
                  "%s in checkpoint doesn't exist in model. Skip.",
                  m_key,
              )
              continue
          if skip_regex and re.findall(skip_regex, m_key):
              logging.info("Skip loading parameter %s.", m_key)
              continue
          if "encoder/class_embedding" in m_key:
              if m_params.shape != model_flat[m_key].shape:
                  model_flat[m_key] = interpolate_class_embedding(
                      model_flat[m_key], m_params
                  )
              else:
                  model_flat[m_key] = m_params
          elif (
              "encoder/VisionTransformer/positional_embedding" in m_key
          ):
              if m_params.shape != model_flat[m_key].shape:
                  model_flat[
                      m_key
                  ] = vivit_model_utils.interpolate_positional_embeddings(
                      m_params, model_flat[m_key].shape[0]
                  )[
                      0
                  ]
              else:
                  model_flat[m_key] = m_params
          elif m_key == "text_encoder/positional_embedding":
              if m_params.shape != model_flat[m_key].shape:
                  cur_len = model_flat[m_key].shape[0]
                  pretrain_len = m_params.shape[0]
                  if pretrain_len > cur_len:
                      model_flat[m_key] = m_params[:cur_len]
                  else:
                      model_flat[m_key] = jnp.concatenate(
                          [m_params, model_flat[m_key][pretrain_len:]],
                          axis=0,
                      )
                  logging.info(
                      "Changing shape of %s from %s to %s.",
                      m_key,
                      m_params.shape,
                      model_flat[m_key].shape,
                  )
              else:
                  model_flat[m_key] = m_params
          elif "conv1" in m_key:
              # backward compatible to 3D conv implementation.
              if m_params.shape != model_flat[m_key].shape:
                  assert (
                      len(m_params.shape) == 5
                      and len(model_flat[m_key].shape) == 2
                  )
                  model_flat[m_key] = jnp.reshape(
                      m_params, model_flat[m_key].shape
                  )
                  logging.info(
                      "Changing shape of %s from %s to %s.",
                      m_key,
                      m_params.shape,
                      model_flat[m_key].shape,
                  )
              else:
                  model_flat[m_key] = m_params
          else:
              logging.info(
                  "Loading %s from checkpoint into model", m_key
              )
              if m_params.shape != model_flat[m_key].shape:
                  raise ValueError(
                      "Inconsistent shapes between the current model (%s) and the "
                      "pretrained model's (%s)."
                      % (model_flat[m_key].shape, m_params.shape)
                  )
              model_flat[m_key] = m_params
      return flax.traverse_util.unflatten_dict(model_flat, sep="/")
  
imports: import collections.abc
  import functools
  import re
  from typing import Any, Callable, Dict, Mapping, Optional, Sequence, Tuple, Union
  from absl import logging
  import flax
  import flax.linen
  import jax
  from jax.example_libraries.optimizers import clip_grads
  import jax.numpy
  import ml_collections
  import optax
  from scenic.common_lib import debug_utils
  from scenic.dataset_lib import dataset_utils
  from scenic.projects.baselines.clip import model
  from scenic.projects.unloc import eval_utils
  from scenic.projects.unloc import model_utils
  from scenic.projects.vivit import model_utils
  from scenic.train_lib import pretrain_utils
  from scenic.train_lib import train_utils
  import scipy.ndimage
label_desc: Use image node_id 3 for calling a global function with example usage: initialize_from_unloc_parameters(params, restored_params, skip_regex) and returns: flax