code: def restore_pretrained_checkpoint(
      checkpoint_path: str,
      train_state: Optional[TrainState] = None,
      assert_exist: bool = False,
      step: Optional[int] = None,
  ) -> TrainState:
      """Restores the last checkpoint.
  
      First restores the checkpoint, which is an instance of TrainState that holds
      the state of training. This function also take care converting pre-Linen
      checkpoints.
  
      Args:
        checkpoint_path: Directory for saving the checkpoint.
        train_state: An instance of TrainState that holds the state of training.
        assert_exist: Assert that there is at least one checkpoint exists in the
          given path.
        step: Step number to load or None to load latest. If specified,
          checkpoint_path must be a directory.
  
      Returns:
        Training state and an int which is the current step.
      """
      if assert_exist:
          glob_path = os.path.join(checkpoint_path, "checkpoint_*")
          if not gfile.glob(glob_path):
              raise ValueError(
                  "No checkpoint for the pretrained model is found in: "
                  f"{checkpoint_path}"
              )
      restored_train_state = checkpoints.restore_checkpoint(
          checkpoint_path, None, step
      )
  
      # no bins at PT case
      if (
          train_state.optimizer.target["text_decoder"][
              "decoder_module"
          ][  # pytype: disable=attribute-error
              "logits_dense"
          ][
              "kernel"
          ].shape
          > restored_train_state["optimizer"]["target"]["text_decoder"][
              "decoder_module"
          ]["logits_dense"]["kernel"].shape
      ):
          x = restored_train_state["optimizer"]["target"][
              "text_decoder"
          ]["decoder_module"]["logits_dense"]["kernel"]
          x2 = train_state.optimizer.target["text_decoder"][
              "decoder_module"
          ][  # pytype: disable=attribute-error
              "logits_dense"
          ][
              "kernel"
          ]
          restored_train_state["optimizer"]["target"]["text_decoder"][
              "decoder_module"
          ]["logits_dense"]["kernel"] = jnp.concatenate(
              [x, x2[:, x.shape[1] :]], 1
          )
          y = restored_train_state["optimizer"]["state"][
              "param_states"
          ]["text_decoder"]["decoder_module"]["logits_dense"]["kernel"]
          y2 = train_state.optimizer.state.param_states[
              "text_decoder"
          ][  # pytype: disable=attribute-error
              "decoder_module"
          ][
              "logits_dense"
          ][
              "kernel"
          ]
          restored_train_state["optimizer"]["state"]["param_states"][
              "text_decoder"
          ]["decoder_module"]["logits_dense"]["kernel"][
              "grad_ema"
          ] = jnp.concatenate(
              [y["grad_ema"], y2.grad_ema[:, y["grad_ema"].shape[1] :]],
              1,
          )
          restored_train_state["optimizer"]["state"]["param_states"][
              "text_decoder"
          ]["decoder_module"]["logits_dense"]["kernel"][
              "grad_sq_ema"
          ] = jnp.concatenate(
              [
                  y["grad_sq_ema"],
                  y2.grad_sq_ema[:, y["grad_sq_ema"].shape[1] :],
              ],
              1,
          )
      # no bins at FT case
      elif (
          train_state.optimizer.target["text_decoder"][
              "decoder_module"
          ][  # pytype: disable=attribute-error
              "logits_dense"
          ][
              "kernel"
          ].shape
          < restored_train_state["optimizer"]["target"]["text_decoder"][
              "decoder_module"
          ]["logits_dense"]["kernel"].shape
      ):
          x = restored_train_state["optimizer"]["target"][
              "text_decoder"
          ]["decoder_module"]["logits_dense"]["kernel"]
          x2 = train_state.optimizer.target["text_decoder"][
              "decoder_module"
          ][  # pytype: disable=attribute-error
              "logits_dense"
          ][
              "kernel"
          ]
          restored_train_state["optimizer"]["target"]["text_decoder"][
              "decoder_module"
          ]["logits_dense"]["kernel"] = x[:, : x2.shape[1]]
          y = restored_train_state["optimizer"]["state"][
              "param_states"
          ]["text_decoder"]["decoder_module"]["logits_dense"]["kernel"]
          y2 = train_state.optimizer.state.param_states[
              "text_decoder"
          ][  # pytype: disable=attribute-error
              "decoder_module"
          ][
              "logits_dense"
          ][
              "kernel"
          ]
          restored_train_state["optimizer"]["state"]["param_states"][
              "text_decoder"
          ]["decoder_module"]["logits_dense"]["kernel"]["grad_ema"] = y[
              "grad_ema"
          ][
              :, : y2.grad_ema.shape[1]
          ]
          restored_train_state["optimizer"]["state"]["param_states"][
              "text_decoder"
          ]["decoder_module"]["logits_dense"]["kernel"][
              "grad_sq_ema"
          ] = y[
              "grad_sq_ema"
          ][
              :, : y2.grad_sq_ema.shape[1]
          ]
  
      # no bins at PT case
      if (
          restored_train_state["optimizer"]["target"][
              "shared_decoder_token_embedder"
          ]["embedding"].shape
          < train_state.optimizer.target[  # pytype: disable=attribute-error
              "shared_decoder_token_embedder"
          ][
              "embedding"
          ].shape
      ):
          x = copy.deepcopy(
              restored_train_state["optimizer"]["state"][
                  "param_states"
              ]["shared_decoder_token_embedder"]["embedding"]
          )
          x2 = copy.deepcopy(
              train_state.optimizer.state.param_states[  # pytype: disable=attribute-error
                  "shared_decoder_token_embedder"
              ][
                  "embedding"
              ]
          )
          restored_train_state["optimizer"]["state"]["param_states"][
              "shared_decoder_token_embedder"
          ] = {
              "embedding": {
                  "grad_ema": jnp.concatenate(
                      [
                          x["grad_ema"],
                          x2.grad_ema[x["grad_ema"].shape[0] :],
                      ],
                      0,
                  ),
                  "grad_sq_ema": jnp.concatenate(
                      [
                          x["grad_sq_ema"],
                          x2.grad_sq_ema[x["grad_sq_ema"].shape[0] :],
                      ],
                      0,
                  ),
              }
          }
          y = copy.deepcopy(
              restored_train_state["optimizer"]["target"][
                  "shared_decoder_token_embedder"
              ]["embedding"]
          )
          y2 = copy.deepcopy(
              train_state.optimizer.target[  # pytype: disable=attribute-error
                  "shared_decoder_token_embedder"
              ][
                  "embedding"
              ]
          )
          restored_train_state["optimizer"]["target"][
              "shared_decoder_token_embedder"
          ] = {"embedding": jnp.concatenate([y, y2[y.shape[0] :]], 0)}
      # no bins at FT case
      elif (
          restored_train_state["optimizer"]["target"][
              "shared_decoder_token_embedder"
          ]["embedding"].shape
          > train_state.optimizer.target[  # pytype: disable=attribute-error
              "shared_decoder_token_embedder"
          ][
              "embedding"
          ].shape
      ):
          x = copy.deepcopy(
              restored_train_state["optimizer"]["state"][
                  "param_states"
              ]["shared_decoder_token_embedder"]["embedding"]
          )
          x2 = copy.deepcopy(
              train_state.optimizer.state.param_states[  # pytype: disable=attribute-error
                  "shared_decoder_token_embedder"
              ][
                  "embedding"
              ]
          )
          restored_train_state["optimizer"]["state"]["param_states"][
              "shared_decoder_token_embedder"
          ] = {
              "embedding": {
                  "grad_ema": x["grad_ema"][: x2.grad_ema.shape[0]],
                  "grad_sq_ema": x["grad_sq_ema"][
                      : x2.grad_sq_ema.shape[0]
                  ],
              }
          }
          y = copy.deepcopy(
              restored_train_state["optimizer"]["target"][
                  "shared_decoder_token_embedder"
              ]["embedding"]
          )
          y2 = copy.deepcopy(
              train_state.optimizer.target[  # pytype: disable=attribute-error
                  "shared_decoder_token_embedder"
              ][
                  "embedding"
              ]
          )
          restored_train_state["optimizer"]["target"][
              "shared_decoder_token_embedder"
          ] = {"embedding": y[: y2.shape[0]]}
  
      if restored_train_state is None:
          raise ValueError(
              "No checkpoint for the pretrained model is found in: "
              f"{checkpoint_path}"
          )
      (
          restored_params,
          restored_model_state,
      ) = get_params_and_model_state_dict(restored_train_state)
      restored_params = flax.core.freeze(restored_params)
      restored_model_state = flax.core.freeze(restored_model_state)
      if train_state:
          new_train_state = train_state
          new_optimizer = train_state.optimizer.replace(
              # Inspect and compare the parameters of the model with the init-model.
              target=inspect_params(
                  expected_params=train_state.optimizer.target,
                  restored_params=restored_params,
                  fail_if_extra=False,
                  fail_if_missing=False,
                  fail_if_shapes_mismatch=False,
              )
          )
      else:
          new_train_state = TrainState()
          new_optimizer = {"target": restored_params}
  
      new_train_state = (
          new_train_state.replace(  # pytype: disable=attribute-error
              optimizer=new_optimizer,
              model_state=restored_model_state,
              global_step=int(restored_train_state["global_step"]),
              rng=restored_train_state["rng"],
              accum_train_time=restored_train_state.get(
                  "accum_train_time", 0
              ),
          )
      )
  
      return new_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 3 for calling a global function with example usage: restore_pretrained_checkpoint(checkpoint_path, train_state, assert_exist, step) and returns: new_train_state