def eval_step(
    train_state: TrainState, batch: Batch, task: Task, rng: jnp.ndarray, *,
    flax_model: nn.Module, tokenizer_dict: Dict[str, TokenizerFn],
    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
    task: Current evaluation task.
    rng: A PRNGKey for current step.
    flax_model: Flax model
    tokenizer_dict: A dictionary of tokenizer and detokenizer
    config: Configurations of the experiment.
    metric_params: Params for metric models.

  Returns:
    Metric features and generated outputs.
  """
  outputs, _ = generate_step(
      train_state,
      batch,
      task,
      rng,
      flax_model=flax_model,
      tokenizer_dict=tokenizer_dict,
      config=config)

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

  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
    task: Current evaluation task.
    rng: A PRNGKey for current step.
    flax_model: Flax model
    tokenizer_dict: A dictionary of tokenizer and detokenizer
    config: Configurations of the experiment.
    metric_params: Params for metric models.

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