code: def mnist_dataset(batch_size):
      (x_train, y_train), _ = tf.keras.datasets.mnist.load_data()
      # The `x` arrays are in uint8 and have values in the range [0, 255].
      # We need to convert them to float32 with values in the range [0, 1]
      x_train = x_train / np.float32(255)
      y_train = y_train.astype(np.int64)
      train_dataset = (
          tf.data.Dataset.from_tensor_slices((x_train, y_train))
          .shuffle(60000)
          .repeat()
          .batch(batch_size)
      )
      return train_dataset
  
file_path: byteps/example/tensorflow/tensorflow2_mnist_bps_MirroredStrategy.py
imports: import tensorflow
  import numpy
  import json
  import os
  import sys
  import argparse
  import byteps.tensorflow
  from byteps.tensorflow.distribute import MirroredStrategy
label_desc: Use image node_id 1 for calling a global function with example usage: mnist_dataset(batch_size) and returns: train_dataset