code: def initialize_from_unloc_train_state(
      train_state: train_utils.TrainState,
      restored_train_state: train_utils.TrainState,
      skip_regex: Optional[str] = None,
  ) -> train_utils.TrainState:
      """Updates UnLoc's train_state with a pretrained UnLoc model weights.
  
      Args:
        train_state: A raw TrainState for the current model.
        restored_train_state: TrainState of the pretrained model.
        skip_regex: Regular expression of parameters to skip loading.
  
      Returns:
        Updated train_state.
      """
  
      params = flax.core.unfreeze(train_state.params)
      restored_params = flax.core.unfreeze(restored_train_state.params)
      params = initialize_from_unloc_parameters(
          params, restored_params, skip_regex
      )
      return train_state.replace(params=flax.core.freeze(params))
  
imports: import collections.abc
  import functools
  import re
  from typing import Any, Callable, Dict, Mapping, Optional, Sequence, Tuple, Union
  from absl import logging
  import flax
  import flax.linen
  import jax
  from jax.example_libraries.optimizers import clip_grads
  import jax.numpy
  import ml_collections
  import optax
  from scenic.common_lib import debug_utils
  from scenic.dataset_lib import dataset_utils
  from scenic.projects.baselines.clip import model
  from scenic.projects.unloc import eval_utils
  from scenic.projects.unloc import model_utils
  from scenic.projects.vivit import model_utils
  from scenic.train_lib import pretrain_utils
  from scenic.train_lib import train_utils
  import scipy.ndimage
label_desc: Use image node_id 4 for calling a global function with example usage: initialize_from_unloc_train_state(train_state, restored_train_state, skip_regex) and returns: train_state