code: def _evaluate_dataset(
      dataset, eval_step_fn, strategy, prediction_writer=None
  ):
      """Evaluates the dataset using the given step_fn and distribution strategy.
  
      Args:
        dataset: A `tf.data.Dataset` of the evaluation data. Each batch contains an
          image tensor and a label tensor.
        eval_step_fn: A callable which takes a batch of (image, label) and returns
          logits. This can have side effects like updating metric objects.
        strategy: A `tf.distribute.Strategy` for running on one or more devices.
        prediction_writer: A file-like object to save each image's prediction for
          further analysis (e.g. union attack).
  
      Returns:
        None.
      """
      dist_dataset = strategy.experimental_distribute_dataset(dataset)
      for step, (dist_images, dist_labels) in enumerate(dist_dataset):
          if step % 50 == 0:
              logging.info("Evaluation step %d", step)
          if strategy:
              dist_logits = strategy.run(
                  eval_step_fn, args=(dist_images, dist_labels)
              )
              logits = strategy.gather(dist_logits, axis=0)
              images = strategy.gather(dist_images, axis=0)
              labels = strategy.gather(dist_labels, axis=0)
          else:
              images, labels = dist_images, dist_labels
              logits = eval_step_fn(images, labels)
          if prediction_writer:
              image_hashes = map(_hash, images)
              predictions = tf.math.argmax(logits, axis=-1).numpy()
              for x_hash, pred, label in zip(
                  image_hashes, predictions, labels.numpy()
              ):
                  prediction_writer.write(
                      f"{x_hash}\t{pred}\t{label}\n"
                  )
  
file_path: neural-structured-learning/research/multi_representation_adversary/multi_representation_adversary/evaluator.py
imports: import functools
  import hashlib
  import os
  import time
  from absl import logging
  import gin
  from multi_representation_adversary import data
  from multi_representation_adversary import helper
  from multi_representation_adversary import resnet
  import tensorflow.compat.v2
label_desc: Use image node_id 2 for calling a global function with example usage: _evaluate_dataset(dataset, eval_step_fn, strategy, prediction_writer) without return types