code: def evaluate_model(
      eval_step_fn: Callable[..., Any],
      state: TrainState,
      batches: Union[Iterable[Example], tf.data.Dataset],
      epoch: int,
      rngs: Optional[Dict[str, Any]] = None,
  ) -> Metrics:
      """Evaluate a model on a dataset."""
      batch_metrics = []
      for i, batch in enumerate(batches):
          batch = batch_to_numpy(batch)
          if rngs is not None:  # New RNG for each step.
              rngs = {
                  name: jax.random.fold_in(rng, i)
                  for name, rng in rngs.items()
              }
  
          metrics = eval_step_fn(state, batch, rngs)
          batch_metrics.append(metrics)
  
      batch_metrics = jax.device_get(batch_metrics)
      metrics = normalize_batch_metrics(batch_metrics)
      logging.info(
          "eval  epoch %03d loss %.4f accuracy %.2f",
          epoch,
          metrics.loss,
          metrics.accuracy * 100,
      )
      return metrics
  
imports: from typing import Any, Callable, Dict, Iterable, Optional, Sequence, Tuple, Union
  from absl import logging
  from flax import struct
  from flax.metrics import tensorboard
  from flax.training import train_state
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  import optax
  import tensorflow
  import input_pipeline
  import models
label_desc: Use image node_id 10 for calling a global function with example usage: evaluate_model(eval_step_fn, state, batches, epoch, rngs) and returns: metrics