import collections


def eval_step(test_batch, snlds_model, num_samples, temperature):
  """Runs evaluation of model on the test set and returns evaluation metrics.

  Args:
    test_batch: a batch of the test data.
    snlds_model: tf.keras.Model, SNLDS model to be evaluated.
    num_samples: int, number of samples per trajectories to use at eval time.
    temperature: float, annealing temperature to use on the model.

  Returns:
    Dictionary of metrics, str -> list[tf.Tensor],
      aggregates the result dictionaries returned by the model.
  """
  test_values = collections.defaultdict(list)
  for _ in range(10):
    result_dict = snlds_model(
        test_batch, temperature, num_samples=num_samples)
    for k, v in result_dict.items():
      test_values[k].append(v)

  return test_values