code: def init_from_pretrain_weights(
      train_state: TrainState,
      restored_params: 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,
  ) -> TrainState:
      """Updates the train_state with pretrained weights.
  
      Args:
        train_state: A raw TrainState for the model.
        restored_params: Loaded parameters.
        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 {}
      model_params = train_state.optimizer.target
      logging.info(
          "model_params: %s",
          jax.tree_map(
              lambda x: x.shape, flax.core.unfreeze(model_params)
          ),
      )
      logging.info(
          "restored_params: %s",
          jax.tree_map(lambda x: x.shape, restored_params),
      )
      model_params = replace_dict(
          model_params,
          restored_params,
          ckpt_prefix_path,
          model_prefix_path,
          name_mapping,
          skip_regex,
          True,
      )
      new_optimizer = train_state.optimizer.replace(target=model_params)
      train_state = (
          train_state.replace(  # pytype: disable=attribute-error
              optimizer=new_optimizer
          )
      )
      return train_state
  
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 1 for calling a global function with example usage: init_from_pretrain_weights(train_state, restored_params, ckpt_prefix_path, model_prefix_path, name_mapping, skip_regex) and returns: train_state