def create_train_state(input_spec: Sequence[Union[Tuple[Tuple[int, ...],
                                                        jnp.dtype],
                                                  Tuple[int, ...], None]],
                       config: ml_collections.ConfigDict, rng: np.ndarray,
                       is_train: bool):
  """Initializes train state and optimizer with a pytree input_spec.

  If the root type of `input_spec` is `Sequence`, each element is fed to the
  model as position arguments. For the video dataset input the expected shape is
  always [local_bs, t, h, w, 3].

  This function interface is different from the standard scenic function
  `initialize_model`, which does three things initialize model, optimizer,
  and train_state.

  Args:
    input_spec: A PyTree whose leaves are (shape, dtype) pairs specifying the
      shape and dtype of the inputs. If unspecified the dtype is the default
      dtype. input_spec =
        [(dataset.meta_data['input_shape'],
          dataset.meta_data.get('input_dtype', dtype))]
    config: Configurations of the initialization.
    rng: Jax rng key.
    is_train: whether is in training mode.

  Returns:
    Model dict,  initial train_state, and learning_rate function.
  """
  model_dict = train_utils.get_vq_model(config)

  rng, init_rng = jax.random.split(rng)
  model_state_dict, init_params_dict = initialize_model(
      model_dict=model_dict,
      input_spec=input_spec,
      config=config,
      rng=init_rng,
      is_train=is_train)

  tx_dict, g_lr_fn, d_lr_fn = get_optimizer(config)
  del d_lr_fn

  g_params = flax.core.unfreeze(init_params_dict['generator'])
  d_params = flax.core.unfreeze(init_params_dict['discriminator'])

  # Create optimizer.
  # We jit this, such that the arrays that are created on the same
  # device as the input is, in this case the CPU. Else they'd be on device[0].
  g_opt_state = jax.jit(
      tx_dict['generator'].init, backend='cpu')(g_params)
  d_opt_state = jax.jit(
      tx_dict['discriminator'].init, backend='cpu')(d_params)

  logging.info('logging generator parameters')
  parameter_overview.log_parameter_overview(g_params)
  logging.info('logging discriminator parameters')
  parameter_overview.log_parameter_overview(d_params)
  if jax.tree_util.tree_leaves(model_state_dict):
    logging.info('logging model states')
    parameter_overview.log_parameter_overview(model_state_dict)
  ema_params = init_params_dict['generator']

  metadata = dict(
      lecam_ema_real=jnp.asarray(0.), lecam_ema_fake=jnp.asarray(0.))

  train_state = TrainState(
      global_step=0,
      g_opt_state=g_opt_state,
      d_opt_state=d_opt_state,
      g_tx=tx_dict['generator'],
      d_tx=tx_dict['discriminator'],
      g_params=g_params,
      d_params=d_params,
      g_model_state=flax.core.freeze(model_state_dict['generator']),
      d_model_state=flax.core.freeze(model_state_dict['discriminator']),
      ema_params=ema_params,
      metadata=metadata,
      rng=rng)
  return model_dict, train_state, g_lr_fn
Initializes train state and optimizer with a pytree input_spec.

  If the root type of `input_spec` is `Sequence`, each element is fed to the
  model as position arguments. For the video dataset input the expected shape is
  always [local_bs, t, h, w, 3].

  This function interface is different from the standard scenic function
  `initialize_model`, which does three things initialize model, optimizer,
  and train_state.

  Args:
    input_spec: A PyTree whose leaves are (shape, dtype) pairs specifying the
      shape and dtype of the inputs. If unspecified the dtype is the default
      dtype. input_spec =
        [(dataset.meta_data['input_shape'],
          dataset.meta_data.get('input_dtype', dtype))]
    config: Configurations of the initialization.
    rng: Jax rng key.
    is_train: whether is in training mode.

  Returns:
    Model dict,  initial train_state, and learning_rate function.
  
videogvt/trainers/vqgan_trainer.py