code: def run_mnist():
      """Run MNIST training and eval loop."""
      mnist_classifier = tf_estimator.Estimator(
          model_fn=model_fn, model_dir=FLAGS.model_dir
      )
  
      # Set up training and evaluation input functions.
      def train_input_fn():
          """Prepare data for training."""
  
          # When choosing shuffle buffer sizes, larger sizes result in better
          # randomness, while smaller sizes use less memory. MNIST is a small
          # enough dataset that we can easily shuffle the full epoch.
          ds = dataset.train(FLAGS.data_dir)
          ds_batched = (
              ds.cache()
              .shuffle(buffer_size=50000)
              .batch(FLAGS.batch_size)
          )
  
          # Iterate through the dataset a set number (`epochs_between_evals`) of times
          # during each training session.
          ds = ds_batched.repeat(FLAGS.epochs_between_evals)
          return ds
  
      def eval_input_fn():
          return (
              dataset.test(FLAGS.data_dir)
              .batch(FLAGS.batch_size)
              .make_one_shot_iterator()
              .get_next()
          )
  
      # Train and evaluate model.
      for _ in range(FLAGS.train_epochs // FLAGS.epochs_between_evals):
          mnist_classifier.train(input_fn=train_input_fn, hooks=None)
          eval_results = mnist_classifier.evaluate(
              input_fn=eval_input_fn
          )
          print("\nEvaluation results:\n\t%s\n" % eval_results)
  
file_path: mesh/examples/mnist.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import mesh_tensorflow
  import mnist_dataset
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
label_desc: Use image node_id 3 for calling a global function with example usage: run_mnist() and returns: ds, dataset