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."""
      name_mapping = name_mapping or {}
  
      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)
      )
  
file_path: scenic/scenic/train_lib/pretrain_utils.py
imports: from collections import abc
  import os
  import re
  from typing import Any, Dict, Mapping, List, Optional, Union
  from absl import logging
  import flax
  from flax.training import checkpoints
  import numpy
  from scenic.train_lib import train_utils
  from tensorflow.io import gfile
label_desc: Use image node_id 1 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