code: def init_from_pretrain_state(
      train_state: train_utils.TrainState,
      pretrain_state: Union[PyTree, train_utils.TrainState],
      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,
  ) -> train_utils.TrainState:
      """Updates the train_state with data from pretrain_state.
  
      Args:
        train_state: A raw TrainState for the model.
        pretrain_state: A TrainState that is loaded with parameters/state of
          a  pretrained model.
        ckpt_prefix_path: Prefix to restored model parameters.
        model_prefix_path: Prefix to the parameters to replace in the subtree model.
        name_mapping: Mapping from parameter names of checkpoint to this model.
        skip_regex: If there is a parameter whose parent keys match the regex,
          the parameter will not be replaced from pretrain_state.
  
      Returns:
        Updated train_state.
      """
      name_mapping = name_mapping or {}
      restored_params = pretrain_state["params"]
      restored_model_state = pretrain_state["model_state"]
      model_params = _replace_dict(
          train_state.params,
          restored_params,
          ckpt_prefix_path,
          model_prefix_path,
          name_mapping,
          skip_regex,
      )
      train_state = train_state.replace(params=model_params)
      # TODO(scenic): Add support for optionally restoring optimizer state.
      if (
          restored_model_state is not None
          and train_state.model_state is not None
          and train_state.model_state
      ):
          if model_prefix_path:
              # Insert model prefix after 'batch_stats'.
              model_prefix_path = ["batch_stats"] + model_prefix_path
              if "batch_stats" in restored_model_state:
                  ckpt_prefix_path = ckpt_prefix_path or []
                  ckpt_prefix_path = ["batch_stats"] + ckpt_prefix_path
          elif (
              "batch_stats" not in restored_model_state
          ):  # Backward compatibility.
              model_prefix_path = ["batch_stats"]
          if ckpt_prefix_path and ckpt_prefix_path[0] != "batch_stats":
              ckpt_prefix_path = ["batch_stats"] + ckpt_prefix_path
          model_state = _replace_dict(
              train_state.model_state,
              restored_model_state,
              ckpt_prefix_path,
              model_prefix_path,
              name_mapping,
              skip_regex,
          )
          train_state = (
              train_state.replace(  # pytype: disable=attribute-error
                  model_state=model_state
              )
          )
      return train_state
  
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 2 for calling a global function with example usage: init_from_pretrain_state(train_state, pretrain_state, ckpt_prefix_path, model_prefix_path, name_mapping, skip_regex) and returns: train_state