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,
  ) -> PyTree:
      """Replaces values in model dictionary with restored ones from checkpoint."""
      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:
              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
          )
          model_flat[m_key] = m_params
  
      return flax.core.freeze(
          flax.traverse_util.unflatten_dict(model_flat)
      )
  
imports: import dataclasses
  import operator
  import re
  from typing import Any, Dict, List, Mapping, Optional, Tuple, Union
  from absl import logging
  from big_vision import utils
  from big_vision.models import vit
  from clu import metric_writers
  import flax
  from flax import jax_utils
  from flax.core import frozen_dict
  import jax
  import jax.numpy
  import ml_collections
  import optax
  from scenic.dataset_lib import dataset_utils
  from scenic.dataset_lib import datasets
  from scenic.projects.knowledge_visual_language.models import constants
  from scenic.projects.t5 import model
  from scenic.train_lib import train_utils
  from scenic.train_lib.train_utils import TrainState
label_desc: Use image node_id 5 for calling a global function with example usage: replace_dict(model, restored, ckpt_prefix_path, model_prefix_path, name_mapping, skip_regex) and returns: flax