def sample_step(train_state: TrainState, batch: Batch, *,
                model_dict: Dict[str, nn.Module],
                config: ml_collections.ConfigDict):
  """Runs a single step to generate samples given input videos.

  Args:
    train_state: The state of training TrainState
    batch: A single batch of data. Dictionary where
      batch['inputs'].shape= device_bs, t, h, w, 3
      batch['batch_mask'].shape= device_bs
      batch['batch_mask'].dtype= float32
      where batch['batch_mask'] > 0 for valid examples
    model_dict: A dictionary of generator and discriminator Flax models
    config: Configurations of the experiment.

  Returns:
    Sampled videos formatted in spatial grids.
  """
  del config
  variables = {
      'params': train_state.g_params,
      **train_state.g_model_state
  }
  ema_variables = {
      'params': train_state.ema_params,
      **train_state.g_model_state
  }

  generator = model_dict['generator']
  generate_fn = functools.partial(generator.apply, variables)
  ema_generate_fn = functools.partial(generator.apply, ema_variables)
  outputs = {
      'original_video': batch['inputs'],
  }
  if 'batch_mask' in batch:
    outputs.update(dict(batch_mask=batch['batch_mask']))

  if 'label' in batch:
    outputs['label'] = batch['label']
  generated_video, _ = generate_fn(batch['inputs'])
  outputs['generated_video'] = jnp.clip(generated_video, 0, 1)
  generated_video_ema, result_dict_ema = ema_generate_fn(batch['inputs'])
  outputs['generated_video_ema'] = jnp.clip(generated_video_ema, 0, 1)
  outputs['generated_tokens_ema'] = result_dict_ema['encoding_indices']
  results = {}
  return outputs, results
Runs a single step to generate samples given input videos.

  Args:
    train_state: The state of training TrainState
    batch: A single batch of data. Dictionary where
      batch['inputs'].shape= device_bs, t, h, w, 3
      batch['batch_mask'].shape= device_bs
      batch['batch_mask'].dtype= float32
      where batch['batch_mask'] > 0 for valid examples
    model_dict: A dictionary of generator and discriminator Flax models
    config: Configurations of the experiment.

  Returns:
    Sampled videos formatted in spatial grids.
  
videogvt/trainers/vqgan_trainer.py