code: def run(model_version: str, pytorch_data_path: str, output_dir: str):
      """Converts VideoMAE checkpoints to Jax, and checks for correctness."""
  
      # First initialise a ViT model and load weights.
      logging.info("Initialising ViT model")
      model_config = get_vivit_config(model_version)
      logging.info(model_config)
      vivit_model = model.ObjectViViT(**model_config)
      init_rngs = {
          "params": random.PRNGKey(0),
          "dropout": random.PRNGKey(1),
      }
  
      rng = random.PRNGKey(0)
      batch, time, height, width, channels = 1, 16, 224, 224, 3
      input_shape = (batch, time, height, width, channels)
      inputs = random.normal(rng, shape=input_shape)
  
      _, params_jax = vivit_model.init_with_output(
          init_rngs, inputs, train=True
      )
      params_jax = flax.core.unfreeze(params_jax["params"])
      params_jax_init = jax.tree_util.tree_map(
          lambda x: x.copy(), params_jax
      )
      # Load VideoMAE model weights.
      logging.info("Loading VideoMAE ViViT weights")
      pretrained_pytorch_data = load_pytorch_data(pytorch_data_path)
      params_pytorch = pretrained_pytorch_data["params"]
  
      # Transfer weights.
      logging.info("Transferring weights")
      params_jax = convert_pytorch_parameters(
          params_pytorch,
          model_config["num_heads"],
          model_config["num_layers"],
      )
      logging.info("Parameter summary:")
      params_jax = flax.core.freeze(params_jax)
      # Check weight transfer was correct.
      if check_weights_different(
          params_jax_init, params_jax, ignore_key_bias=True
      ):
          logging.info("All model weights changed from initialisation.")
  
      # Save Scenic train state.
      logging.info("Saving converted checkpoint to %s", output_dir)
      train_state = train_utils.TrainState(
          global_step=0, params=params_jax, model_state={}
      )
      train_utils.save_checkpoint(
          output_dir, train_state, overwrite=True
      )
  
imports: import pickle
  from typing import Any, Dict, Union
  from absl import app
  from absl import flags
  from absl import logging
  import flax
  import jax
  from jax import random
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.projects.objectvivit import model
  from scenic.train_lib import train_utils
label_desc: Use image node_id 8 for calling a global function with example usage: run(model_version, pytorch_data_path, output_dir) without return types