code: def create_train_state(ckpt_path, config):
      """Gets the train state from the input ckpt.
  
      Args:
        ckpt_path: Path to the ckpt to be loaded..
        config: The configuration used to train the model.
  
      Returns:
        Loaded train state and a function to return the logits.
      """
      ckpt_info = ckpt_path.split("/")
      ckpt_dir = "/".join(ckpt_info[:-1])
      ckpt_num = ckpt_info[-1].split("_")[-1]
  
      rng = jax.random.PRNGKey(config.rng_seed)
      data_rng, rng = jax.random.split(rng)
      dataset = train_utils.get_dataset(
          config, data_rng, dataset_service_address=None
      )
      model = matvit.MatViTMultiLabelClassificationModel(
          config, dataset.meta_data
      )
  
      init_rng, _ = jax.random.split(rng)
      (params, model_state, _, _) = train_utils.initialize_model(
          model_def=model.flax_model,
          input_spec=[
              (
                  dataset.meta_data["input_shape"],
                  dataset.meta_data.get("input_dtype", jnp.float32),
              )
          ],
          config=config,
          rngs=init_rng,
      )
      train_state = train_utils.TrainState(
          params=params, model_state=model_state
      )
      train_state, _ = train_utils.restore_checkpoint(
          ckpt_dir,
          train_state,
          assert_exist=True,
          step=int(ckpt_num),
      )
      if not _MATVIT_DIMS.value:
          matvit_dims = None
      else:
          matvit_dims = [
              int(val) for val in _MATVIT_DIMS.value.split(",")
          ]
          assert (
              len(matvit_dims) == config.model.num_layers
          ), "Number of matvit dimensions needs to match the number of layers."
  
      logits_fn = functools.partial(
          get_logits_fn, flax_model=model.flax_model, dims=matvit_dims
      )
      return train_state, logits_fn
  
file_path: scenic/scenic/projects/matvit/classification_eval_main.py
imports: import functools
  from typing import Callable, Dict, List, Optional, Tuple
  from absl import app
  from absl import flags
  from absl import logging
  from flax import jax_utils
  import flax.linen
  import jax
  import jax.numpy
  import jax.profiler
  import ml_collections
  import numpy
  from scenic.projects.matvit import matvit
  from scenic.train_lib import train_utils
label_desc: Use image node_id 3 for calling a global function with example usage: create_train_state(ckpt_path, config) and returns: train_state,  logits_fn