code: def init_state(
      model: base_model.BaseModel,
      dataset: dataset_utils.Dataset,
      config: ml_collections.ConfigDict,
      workdir: str,
      rng: jnp.ndarray,
  ):
      """Initialize the train state."""
  
      encoder_input_shape = dataset.meta_data["encoder_input_shape"]
      encoder_input_dtype = dataset.meta_data.get(
          "encoder_input_dtype", jnp.float32
      )
      encoder_input_text_dtype = dataset.meta_data.get(
          "encoder_input_text_dtype", jnp.int32
      )
      decoder_input_shape = dataset.meta_data["decoder_input_shape"]
      decoder_input_dtype = dataset.meta_data.get(
          "decoder_input_dtype", jnp.int32
      )
  
      encoder_input_spec = {}
      if isinstance(encoder_input_shape, dict):
          for mod in config.dataset_configs.modalities:
              mod_spec = None
              if mod in encoder_input_shape:
                  logging.info(
                      "Modality %s is present for this dataset", mod
                  )
                  local_encoder_input_dtype = encoder_input_dtype
                  if mod == "text":
                      local_encoder_input_dtype = (
                          encoder_input_text_dtype
                      )
                  mod_spec = (
                      encoder_input_shape[mod],
                      local_encoder_input_dtype,
                  )
              encoder_input_spec[mod] = mod_spec
      else:
          encoder_input_spec = (
              encoder_input_shape,
              encoder_input_dtype,
          )
  
      decoder_input_spec = {
          k: (v, decoder_input_dtype)
          for k, v in decoder_input_shape.items()
      }
  
      # Initialize model.
      rng, init_rng = jax.random.split(rng)
      (
          params,
          model_state,
          num_params,
          gflops,
      ) = train_utils.initialize_model_with_pytree(
          model_def=model.flax_model,
          input_spec=(encoder_input_spec, decoder_input_spec),
          config=config,
          rngs=init_rng,
      )
      logging.info(
          "The model has %d params, uses %d gflops",
          num_params,
          gflops or -1,
      )
  
      # Create the optimizer.
      # We jit this, such that the arrays that are created are created on the same
      # device as the input is, in this case the CPU. Else they'd be on device[0].
      optimizer = jax.jit(
          optimizers.get_optimizer(config).create, backend="cpu"
      )(params)
      del params  # Do not keep a copy of the initial params.
  
      rng, train_rng = jax.random.split(rng)
      train_state = train_utils.TrainState(
          global_step=0,
          optimizer=optimizer,
          model_state=model_state,
          rng=train_rng,
          accum_train_time=0,
      )
      start_step = train_state.global_step
      if config.checkpoint:
          logging.info("Continuing training from the checkpoint")
          train_state, params_axes = vid2seq_train_utils.pop_axes_names(
              train_state, "params_axes"
          )
          train_state, start_step = train_utils.restore_checkpoint(
              workdir, train_state
          )
          train_state = vid2seq_train_utils.re_add_axis_names(
              train_state, params_axes, "params_axes"
          )
  
      if start_step == 0 and config.get("init_from"):
          if config.init_from.get("checkpoint_path"):
              restored_model_cfg = config.init_from.get("model_config")
              init_checkpoint_path = config.init_from.get(
                  "checkpoint_path"
              )
              step = config.init_from.get("step")
              restored_train_state = (
                  pretrain_utils.restore_pretrained_checkpoint(
                      init_checkpoint_path,
                      train_state,
                      assert_exist=True,
                      step=step,
                  )
              )
  
              # init T5 encoder if in model but not in ckpt
              if (
                  config.model.encoder.encoder_type
                  in ["t5_encoder", "cat_encoder"]
                  and "text" in config.dataset_configs.modalities
              ):
                  if (
                      "video_encoder"
                      not in restored_train_state.optimizer.target[
                          "encoder"
                      ]
                      and "video_encoder"
                      in train_state.optimizer.target["encoder"]
                  ):
                      train_state = load_encoder_params(
                          train_state, config
                      )
                      x = unfreeze(
                          restored_train_state.optimizer.target
                      )
                      x["encoder"] = copy.deepcopy(
                          train_state.optimizer.target["encoder"]
                      )
                      optimizer = (
                          restored_train_state.optimizer.replace(
                              target=x
                          )
                      )
                      restored_train_state = (
                          restored_train_state.replace(
                              optimizer=optimizer
                          )
                      )
  
                      y = unfreeze(
                          restored_train_state.optimizer.state.param_states
                      )
                      y["encoder"]["video_encoder"] = copy.deepcopy(
                          train_state.optimizer.state.param_states[
                              "encoder"
                          ]["video_encoder"]
                      )
                      state = (
                          restored_train_state.optimizer.state.replace(
                              param_states=y
                          )
                      )
                      optimizer = (
                          restored_train_state.optimizer.replace(
                              state=state
                          )
                      )
                      restored_train_state = (
                          restored_train_state.replace(
                              optimizer=optimizer
                          )
                      )
  
              # throw away T5 encoder if in checkpoint but not in model
              if (
                  config.model.encoder.encoder_type
                  in ["t5_encoder", "cat_encoder"]
                  and "text" not in config.dataset_configs.modalities
              ):
                  if (
                      "video_encoder"
                      in restored_train_state.optimizer.target[
                          "encoder"
                      ]
                      and "video_encoder"
                      not in train_state.optimizer.target["encoder"]
                  ):
                      x = unfreeze(
                          restored_train_state.optimizer.target
                      )
                      if (
                          "encoder" in x
                          and "video_encoder" in x["encoder"]
                      ):
                          del x["encoder"]["video_encoder"]
                      optimizer = (
                          restored_train_state.optimizer.replace(
                              target=x
                          )
                      )
                      restored_train_state = (
                          restored_train_state.replace(
                              optimizer=optimizer
                          )
                      )
  
                      y = unfreeze(
                          restored_train_state.optimizer.state.param_states
                      )
                      if (
                          "encoder" in y
                          and "video_encoder" in y["encoder"]
                      ):
                          del y["encoder"]["video_encoder"]
                      state = (
                          restored_train_state.optimizer.state.replace(
                              param_states=y
                          )
                      )
                      optimizer = (
                          restored_train_state.optimizer.replace(
                              state=state
                          )
                      )
                      restored_train_state = (
                          restored_train_state.replace(
                              optimizer=optimizer
                          )
                      )
  
              train_state = load_utils.initialise_from_train_state(
                  config,
                  train_state,
                  restored_train_state,
                  restored_model_cfg,
                  restore_output_proj=False,
              )
          else:
              # Seperately intiialize encoders and decoders
              if config.model.encoder.encoder_type in [
                  "t5_encoder",
                  "cat_encoder",
              ]:
                  train_state = load_encoder_params(train_state, config)
              train_state = load_decoder_params(train_state, config)
  
      elif start_step == 0:
          logging.info(
              "Training completely from scratch. "
              "Not restoring from any checkpoint."
          )
      return train_state, start_step
  
imports: import copy
  import dataclasses
  import functools
  import json
  import os
  import random
  from typing import Any, Callable, Dict, List, Optional, Tuple
  from absl import logging
  from clu import metric_writers
  from clu import periodic_actions
  from dmvr import tokenizers
  from flax import jax_utils
  from flax.core import unfreeze
  import flax.linen
  import jax
  import jax.example_libraries.optimizers
  import jax.numpy
  import jax.profiler
  import ml_collections
  import numpy
  from scenic.dataset_lib import dataset_utils
  from scenic.model_lib.base_models import base_model
  from scenic.projects.t5 import model
  from scenic.projects.t5 import tokenizer
  from scenic.projects.vid2seq import load_utils
  from scenic.projects.vid2seq import models
  from scenic.projects.vid2seq import train_utils
  from scenic.projects.vid2seq import dvc_eval
  from scenic.train_lib_deprecated import lr_schedules
  from scenic.train_lib_deprecated import optimizers
  from scenic.train_lib_deprecated import pretrain_utils
  from scenic.train_lib_deprecated import train_utils
label_desc: Use image node_id 9 for calling a global function with example usage: init_state(model, dataset, config, workdir, rng) and returns: train_state,  start_step