code: def serving_call_mnist(images):
      """Send an RPC or REST request to the model server.
  
      Args:
        images: A numpy.ndarray of shape [B, 28, 28, 1] with the batch of images to
          perform inference on.
  
      Returns:
        A numpy.ndarray of shape [B, 10] with the one-hot inference response.
      """
      if _USE_GRPC.value:
          channel = grpc.insecure_channel(
              _PREDICTION_SERVICE_ADDR.value
          )
          stub = prediction_service_pb2_grpc.PredictionServiceStub(
              channel
          )
  
          request = predict_pb2.PredictRequest()
          request.model_spec.name = _MODEL_SPEC_NAME.value
          request.model_spec.signature_name = (
              tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY
          )
          # You can see the name of the input ("inputs") in the SavedModel dump.
          request.inputs["inputs"].CopyFrom(
              tf.make_tensor_proto(
                  images, dtype=images.dtype, shape=images.shape
              )
          )
          response = stub.Predict(request)
          # We could also use response.outputs["output_0"], where "output_0" is the
          # name of the output (which you can see in the SavedModel dump.)
          # Alternatively, we just get the first output.
          (outputs,) = response.outputs.values()
          return tf.make_ndarray(outputs)
      else:
          # Use the HTTP REST api
          images_json = json.dumps(images.tolist())
          # You can see the name of the input ("inputs") in the SavedModel dump.
          data = f'{{"inputs": {images_json}}}'
          predict_url = f"http://{_PREDICTION_SERVICE_ADDR.value}/v1/models/{_MODEL_SPEC_NAME.value}:predict"
          response = requests.post(predict_url, data=data)
          if response.status_code != 200:
              msg = (
                  f"Received error response {response.status_code} from model "
                  f"server: {response.text}"
              )
              raise ValueError(msg)
          outputs = response.json()["outputs"]
          return np.array(outputs)
  
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 1 for calling a global function with example usage: serving_call_mnist(images) and returns: tf, np