code: def main(_):
      logging.info("Loading the MNIST TensorFlow dataset")
      train_ds = mnist_lib.load_mnist(
          tfds.Split.TRAIN, batch_size=mnist_lib.train_batch_size
      )
      test_ds = mnist_lib.load_mnist(
          tfds.Split.TEST, batch_size=_SERVING_BATCH_SIZE
      )
  
      (flax_predict, flax_params) = mnist_lib.FlaxMNIST.train(
          train_ds, test_ds, _NUM_EPOCHS.value
      )
  
      def predict(image):
          return flax_predict(flax_params, image)
  
      # Convert Flax model to TF function.
      tf_predict = tf.function(
          jax2tf.convert(predict, enable_xla=False),
          input_signature=[
              tf.TensorSpec(
                  shape=[_SERVING_BATCH_SIZE, 28, 28, 1],
                  dtype=tf.float32,
                  name="input",
              )
          ],
          autograph=False,
      )
  
      # Convert TF function to TF Lite format.
      converter = tf.lite.TFLiteConverter.from_concrete_functions(
          [tf_predict.get_concrete_function()], tf_predict
      )
      converter.target_spec.supported_ops = [
          tf.lite.OpsSet.TFLITE_BUILTINS,  # enable TensorFlow Lite ops.
          tf.lite.OpsSet.SELECT_TF_OPS,  # enable TensorFlow ops.
      ]
      tflite_float_model = converter.convert()
  
      # Show model size in KBs.
      float_model_size = len(tflite_float_model) / 1024
      print("Float model size = %dKBs." % float_model_size)
  
      # Re-convert the model to TF Lite using quantization.
      converter.optimizations = [tf.lite.Optimize.DEFAULT]
      tflite_quantized_model = converter.convert()
  
      # Show model size in KBs.
      quantized_model_size = len(tflite_quantized_model) / 1024
      print("Quantized model size = %dKBs," % quantized_model_size)
      print(
          "which is about %d%% of the float model size."
          % (quantized_model_size * 100 / float_model_size)
      )
  
      # Evaluate the TF Lite float model. You'll find that its accuracy is identical
      # to the original Flax model because they are essentially the same model
      # stored in different format.
      float_accuracy = evaluate_tflite_model(
          tflite_float_model, test_ds
      )
      print("Float model accuracy = %.4f" % float_accuracy)
  
      # Evalualte the TF Lite quantized model.
      # Don't be surprised if you see quantized model accuracy is higher than
      # the original float model. It happens sometimes :)
      quantized_accuracy = evaluate_tflite_model(
          tflite_quantized_model, test_ds
      )
      print("Quantized model accuracy = %.4f" % quantized_accuracy)
      print(
          "Accuracy drop = %.4f" % (float_accuracy - quantized_accuracy)
      )
  
      f = open(_TFLITE_FILE_PATH.value, "wb")
      f.write(tflite_quantized_model)
      f.close()
  
file_path: jax/jax/experimental/jax2tf/examples/tflite/mnist/mnist.py
imports: import logging
  from absl import app
  from absl import flags
  from jax.experimental import jax2tf
  from jax.experimental.jax2tf.examples import mnist_lib
  import numpy
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 2 for calling a global function with example usage: main(_) and returns: flax_predict