code: def init_from_unloc_checkpoint(
      config: ml_collections.ConfigDict,
      train_state: train_utils.TrainState,
  ) -> train_utils.TrainState:
      """Initialize train state from an UnLoc checkpoint.
  
      The checkpoint can be specified either by an xid:
      config.init_from.xm = (55208837, 1)
      or a file path:
      config.init_from.checkpoint_path = '/is-d/home/foo/55208837/1'
  
      Args:
        config: Config points to checkpoint location.
        train_state: TrainState of currement model.
  
      Returns:
        Updated train_state.
      Raises:
        RuntimeError: if checkpoint is not provided by the config.
      """
      init_checkpoint_path = None
      checkpoint_path = config.init_from.get("checkpoint_path")
      if checkpoint_path is not None:
          init_checkpoint_path = checkpoint_path
      if init_checkpoint_path is None:
          raise RuntimeError(
              'Set either "xm" or a file path to UnLoc checkpoint.'
          )
  
      restored_train_state = (
          pretrain_utils.restore_pretrained_checkpoint(
              init_checkpoint_path, train_state, assert_exist=True
          )
      )
      return initialize_from_unloc_train_state(
          train_state,
          restored_train_state,
          skip_regex=config.init_from.get("skip_regex"),
      )
  
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 5 for calling a global function with example usage: init_from_unloc_checkpoint(config, train_state) and returns: initialize_from_unloc_train_state