def init_from_pretrained_checkpoint(
    checkpoint_path: str,
    train_state: TrainState,
    expand_embedding: bool = False,
    config: Optional[ml_collections.ConfigDict] = None):
  """Initialize the train state with a pretrained checkpoint.

  First restores the checkpoint, then expands position embedding if needed.

  Args:
    checkpoint_path: Directory for saving the checkpoint.
    train_state: An instance of TrainState that holds the state of training.
    expand_embedding: Whether expand embedding for condition tokens.
    config: Current model config.

  Returns:
    Training state and an int which is the current step.
  """
  restored_train_state = checkpoints.restore_checkpoint(checkpoint_path, None)
  if restored_train_state is None:
    raise ValueError('No checkpoint for the pretrained model is found in: '
                     f'{checkpoint_path}')
  restored_train_state = flax.core.freeze(restored_train_state)
  if 'optimizer' in restored_train_state:
    # Legacy train_state
    restored_params = restored_train_state['optimizer']['target']
  else:
    restored_params = restored_train_state['params']
  restored_model_state = restored_train_state.get('model_state')
  if expand_embedding:
    assert config is not None, ('config must be provided when '
                                'expand_embedding=True')
    latent_seq_len = get_latent_seq_len(config)
    vq_codebook_size = get_vq_codebook_size(config)
    num_classes = get_num_classes(config)
    restored_params = flax.core.unfreeze(restored_params)
    restored_params = expand_pos_embedding(
        restored_params, train_state.params, latent_seq_len=latent_seq_len)
    restored_params = expand_word_embedding(
        restored_params, train_state.params, vq_codebook_size=vq_codebook_size,
        num_classes=num_classes)
    restored_params = flax.core.freeze(restored_params)
  # pytype: disable=attribute-error
  new_train_state = train_state.replace(
      params=restored_params, model_state=restored_model_state)
  # pytype: enable=attribute-error
  return new_train_state
Initialize the train state with a pretrained checkpoint.

  First restores the checkpoint, then expands position embedding if needed.

  Args:
    checkpoint_path: Directory for saving the checkpoint.
    train_state: An instance of TrainState that holds the state of training.
    expand_embedding: Whether expand embedding for condition tokens.
    config: Current model config.

  Returns:
    Training state and an int which is the current step.
  
videogvt/train_lib/maskgvt_train_lib.py