def eval_step(
    train_state: TrainState,
    batch: Batch,
    *,
    model_dict: Dict[str, nn.Module],
    config: ml_collections.ConfigDict,
    metric_params: Dict[str, Any]) -> Tuple[Dict[str, Any], Dict[str, Any]]:
  """Runs a single step of evaluation.

  Args:
    train_state: The state of evaluating TrainState
    batch: A single batch of data. Dictionary where batch['inputs'].shape=
      device_bs, t, h, w, 3
    model_dict: A dictionary of generator and discriminator Flax models
    config: Configurations of the experiment.
    metric_params: Params for metric models.

  Returns:
    Metric features and generated outputs.
  """
  outputs, _ = sample_step(
      train_state, batch, model_dict=model_dict, config=config)

  dataset_type = config.get('dataset_type', 'video')
  if dataset_type == 'image':
    # Extend to 16 frames for IS and FVD.
    batch_mask = outputs.pop('batch_mask', None)
    tokens = outputs.pop('generated_tokens_ema')
    outputs = jax.tree_util.tree_map(lambda x: jnp.tile(x, (1, 16, 1, 1, 1)),
                                     outputs)
    outputs['generated_tokens_ema'] = tokens
    if batch_mask is not None:
      outputs['batch_mask'] = batch_mask

  features = eval_utils.eval_step_get_features(
      outputs,
      metric_params=metric_params,
      model_suffix_list=['', '_ema'],
      config=config)

  if dataset_type == 'image':
    # Reduce back to 1 frame.
    batch_mask = outputs.pop('batch_mask', None)
    tokens = outputs.pop('generated_tokens_ema')
    outputs = jax.tree_util.tree_map(lambda x: x[:, :1], outputs)
    outputs['generated_tokens_ema'] = tokens
    if batch_mask is not None:
      outputs['batch_mask'] = batch_mask

  if config.eval.get('enable_lpips', False):
    outputs = jax.lax.all_gather(outputs, axis_name='device', tiled=True)
  return features, outputs
Runs a single step of evaluation.

  Args:
    train_state: The state of evaluating TrainState
    batch: A single batch of data. Dictionary where batch['inputs'].shape=
      device_bs, t, h, w, 3
    model_dict: A dictionary of generator and discriminator Flax models
    config: Configurations of the experiment.
    metric_params: Params for metric models.

  Returns:
    Metric features and generated outputs.
  
videogvt/trainers/vqgan_trainer.py