code: def train_and_save():
      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=mnist_lib.test_batch_size
      )
  
      if SHOW_IMAGES.value:
          mnist_lib.plot_images(
              train_ds, 1, 5, "Training images", inference_fn=None
          )
  
      the_model_class = pick_model_class()
      model_dir = savedmodel_dir(with_version=True)
  
      if _GENERATE_MODEL.value:
          model_descr = model_description()
          logging.info("Generating model for %s", model_descr)
          (predict_fn, predict_params) = the_model_class.train(
              train_ds,
              test_ds,
              num_epochs=_NUM_EPOCHS.value,
              with_classifier=_MODEL_CLASSIFIER_LAYER.value,
          )
  
          if _SERVING_BATCH_SIZE.value == -1:
              # Batch-polymorphic SavedModel
              input_signatures = [
                  tf.TensorSpec(
                      (None,) + mnist_lib.input_shape, tf.float32
                  ),
              ]
              polymorphic_shapes = "(batch, ...)"
          else:
              input_signatures = [
                  # The first one will be the serving signature
                  tf.TensorSpec(
                      (_SERVING_BATCH_SIZE.value,)
                      + mnist_lib.input_shape,
                      tf.float32,
                  ),
                  tf.TensorSpec(
                      (mnist_lib.train_batch_size,)
                      + mnist_lib.input_shape,
                      tf.float32,
                  ),
                  tf.TensorSpec(
                      (mnist_lib.test_batch_size,)
                      + mnist_lib.input_shape,
                      tf.float32,
                  ),
              ]
              polymorphic_shapes = None
  
          logging.info("Saving model for %s", model_descr)
          saved_model_lib.convert_and_save_model(
              predict_fn,
              predict_params,
              model_dir,
              with_gradient=True,
              input_signatures=input_signatures,
              polymorphic_shapes=polymorphic_shapes,
              compile_model=_COMPILE_MODEL.value,
          )
  
          if _TEST_SAVEDMODEL.value:
              (
                  tf_accelerator,
                  tolerances,
              ) = tf_accelerator_and_tolerances()
              with tf.device(tf_accelerator):
                  logging.info("Testing savedmodel")
                  pure_restored_model = tf.saved_model.load(model_dir)
  
                  if (
                      SHOW_IMAGES.value
                      and _MODEL_CLASSIFIER_LAYER.value
                  ):
                      mnist_lib.plot_images(
                          test_ds,
                          1,
                          5,
                          f"Inference results for {model_descr}",
                          inference_fn=pure_restored_model,
                      )
  
                  test_input = np.ones(
                      (mnist_lib.test_batch_size,)
                      + mnist_lib.input_shape,
                      dtype=np.float32,
                  )
                  np.testing.assert_allclose(
                      pure_restored_model(
                          tf.convert_to_tensor(test_input)
                      ),
                      predict_fn(predict_params, test_input),
                      **tolerances,
                  )
  
      if _SHOW_MODEL.value:
  
          def print_model(model_dir: str):
              cmd = f"saved_model_cli show --all --dir {model_dir}"
              print(cmd)
              os.system(cmd)
  
          print_model(model_dir)
  
file_path: jax/jax/experimental/jax2tf/examples/saved_model_main.py
imports: import logging
  import os
  from absl import app
  from absl import flags
  from jax.experimental.jax2tf.examples import mnist_lib
  from jax.experimental.jax2tf.examples import saved_model_lib
  import numpy
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 1 for calling a global function with example usage: train_and_save() without return types