code: def main(_):
      if _COUNT_IMAGES.value % _SERVING_BATCH_SIZE.value != 0:
          raise ValueError(
              f"The count_images ({_COUNT_IMAGES.value}) must be a "
              "multiple of "
              f"serving_batch_size ({_SERVING_BATCH_SIZE.value})"
          )
      test_ds = mnist_lib.load_mnist(
          tfds.Split.TEST, batch_size=_SERVING_BATCH_SIZE.value
      )
      images_and_labels = tfds.as_numpy(
          test_ds.take(_COUNT_IMAGES.value // _SERVING_BATCH_SIZE.value)
      )
  
      accurate_count = 0
      for batch_idx, (images, labels) in enumerate(images_and_labels):
          predictions_one_hot = serving_call_mnist(images)
          predictions_digit = np.argmax(predictions_one_hot, axis=1)
          labels_digit = np.argmax(labels, axis=1)
          accurate_count += np.sum(labels_digit == predictions_digit)
          running_accuracy = (
              100.0
              * accurate_count
              / (1 + batch_idx)
              / _SERVING_BATCH_SIZE.value
          )
          logging.info(
              " predicted digits = %s labels %s. Running accuracy %.3f%%",
              predictions_digit,
              labels_digit,
              running_accuracy,
          )
  
file_path: jax/jax/experimental/jax2tf/examples/serving/model_server_request.py
imports: import grpc
  import json
  import logging
  import requests
  from absl import app
  from absl import flags
  from jax.experimental.jax2tf.examples import mnist_lib
  import numpy
  import tensorflow
  import tensorflow_datasets
  from tensorflow_serving.apis import predict_pb2
  from tensorflow_serving.apis import prediction_service_pb2_grpc
label_desc: Use image node_id 2 for calling a global function with example usage: main(_) without return types