code: def replace_dict(
      model: PyTree,
      restored: PyTree,
      ckpt_prefix_path: Optional[List[str]] = None,
      model_prefix_path: Optional[List[str]] = None,
      name_mapping: Optional[Mapping[str, str]] = None,
      skip_regex: Optional[str] = None,
      compat: bool = False,
  ) -> PyTree:
      """Restores checkpoint in model dictionary, compatible with additional timestamp tokens."""
  
      model = flax.core.unfreeze(
          model
      )  # pytype: disable=wrong-arg-types
      restored = flax.core.unfreeze(
          restored
      )  # pytype: disable=wrong-arg-types
  
      if ckpt_prefix_path:
          for p in ckpt_prefix_path:
              restored = restored[p]
  
      if model_prefix_path:
          for p in reversed(model_prefix_path):
              restored = {p: restored}
  
      # Flatten nested parameters to a dict of str -> tensor. Keys are tuples
      # from the path in the nested dictionary to the specific tensor. E.g.,
      # {'a1': {'b1': t1, 'b2': t2}, 'a2': t3}
      # -> {('a1', 'b1'): t1, ('a1', 'b2'): t2, ('a2',): t3}.
      restored_flat = flax.traverse_util.flatten_dict(
          dict(restored), keep_empty_nodes=True
      )
      model_flat = flax.traverse_util.flatten_dict(
          dict(model), keep_empty_nodes=True
      )
  
      for m_key, m_params in restored_flat.items():
          # pytype: disable=attribute-error
          for name, to_replace in name_mapping.items():
              m_key = tuple(
                  to_replace if k == name else k for k in m_key
              )
          # pytype: enable=attribute-error
          m_key_str = "/".join(m_key)
          if m_key not in model_flat:
              if ("tmp_" + m_key[0], *m_key[1:]) in model_flat:
                  m_key = ("tmp_" + m_key[0], *m_key[1:])
              else:
                  logging.warning(
                      "%s in checkpoint doesn't exist in model. Skip.",
                      m_key_str,
                  )
                  continue
          if skip_regex and re.findall(skip_regex, m_key_str):
              logging.info("Skip loading parameter %s.", m_key_str)
              continue
          logging.info(
              "Loading %s from checkpoint into model", m_key_str
          )
          # shared_decoder_token_embedder/embedding 32128 768
          if compat and m_params.shape[0] < model_flat[m_key].shape[0]:
              logging.info(
                  "Loading everything but the last axis 0 %d values",
                  (model_flat[m_key].shape[0] - m_params.shape[0]),
              )
              model_flat[m_key] = jnp.concatenate(
                  [
                      m_params,
                      model_flat[m_key][
                          -(
                              model_flat[m_key].shape[0]
                              - m_params.shape[0]
                          ) :
                      ],
                  ],
                  axis=0,
              )
          # text_decoder/decoder_module/logits_dense/kernel 768 32128
          elif (
              compat and m_params.shape[0] > model_flat[m_key].shape[0]
          ):  # tmp only
              continue
          elif (
              compat
              and len(m_params.shape) == 2
              and m_params.shape[1] < model_flat[m_key].shape[1]
          ):
              logging.info(
                  "Loading everything but the last axis 1 %d values",
                  (model_flat[m_key].shape[1] - m_params.shape[1]),
              )
              model_flat[m_key] = jnp.concatenate(
                  [
                      m_params,
                      model_flat[m_key][
                          :,
                          -(
                              model_flat[m_key].shape[1]
                              - m_params.shape[1]
                          ) :,
                      ],
                  ],
                  axis=1,
              )
          elif (
              compat
              and len(m_params.shape) == 2
              and m_params.shape[1] > model_flat[m_key].shape[1]
          ):
              continue  # tmp only
          else:
              model_flat[m_key] = m_params
  
      return flax.core.freeze(
          flax.traverse_util.unflatten_dict(model_flat)
      )
  
file_path: scenic/scenic/projects/vid2seq/load_utils.py
imports: import copy
  import os
  import re
  from typing import List, Mapping, Optional, Union, Any
  from absl import logging
  import flax
  from flax.training import checkpoints
  import jax
  import jax.numpy
  import ml_collections
  from scenic.common_lib import debug_utils
  from scenic.train_lib_deprecated.pretrain_utils import get_params_and_model_state_dict
  from scenic.train_lib_deprecated.pretrain_utils import inspect_params
  from scenic.train_lib_deprecated.train_utils import TrainState
  from tensorflow.io import gfile
label_desc: Use image node_id 2 for calling a global function with example usage: replace_dict(model, restored, ckpt_prefix_path, model_prefix_path, name_mapping, skip_regex, compat) and returns: flax