code: def main(max_epoch):
      sc = init_nncontext()
  
      training_rdd = get_data_rdd("train", sc)
      testing_rdd = get_data_rdd("test", sc)
  
      dataset = TFDataset.from_rdd(
          training_rdd,
          features=(tf.float32, [28, 28, 1]),
          labels=(tf.int32, []),
          batch_size=320,
          val_rdd=testing_rdd,
      )
  
      model = tf.keras.Sequential(
          [
              tf.keras.layers.Flatten(input_shape=(28, 28, 1)),
              tf.keras.layers.Dense(64, activation="relu"),
              tf.keras.layers.Dense(64, activation="relu"),
              tf.keras.layers.Dense(10, activation="softmax"),
          ]
      )
  
      model.compile(
          optimizer=tf.keras.optimizers.RMSprop(),
          loss="sparse_categorical_crossentropy",
          metrics=["accuracy"],
      )
  
      keras_model = KerasModel(model)
  
      keras_model.fit(dataset, epochs=max_epoch, distributed=True)
  
      eval_dataset = TFDataset.from_rdd(
          testing_rdd,
          features=(tf.float32, [28, 28, 1]),
          labels=(tf.int32, []),
          batch_per_thread=80,
      )
      result = keras_model.evaluate(eval_dataset)
  
      print(result)
      # >> [0.08865142822265625, 0.9722]
  
      # the following assert is used for internal testing
      assert result["acc Top1Accuracy"] > 0.95
  
      model.save_weights("/tmp/mnist_keras.h5")
  
imports: import sys
  import tensorflow
  import numpy
  from zoo import init_nncontext
  from zoo.tfpark import KerasModel, TFDataset
label_desc: Use image node_id 2 for calling a global function with example usage: main(max_epoch) without return types