code: def get_preprocessed_mnist_data(
      img_rows=28, img_cols=28, num_classes=10, is_quantized_model=False
  ):
      """Get data for mnist training and evaluation."""
      (x_train, y_train), (
          x_test,
          y_test,
      ) = tf.keras.datasets.mnist.load_data()
  
      if tf.keras.backend.image_data_format() == "channels_first":
          x_train = x_train.reshape(
              x_train.shape[0], 1, img_rows, img_cols
          )
          x_test = x_test.reshape(
              x_test.shape[0], 1, img_rows, img_cols
          )
          input_shape = (1, img_rows, img_cols)
      else:
          x_train = x_train.reshape(
              x_train.shape[0], img_rows, img_cols, 1
          )
          x_test = x_test.reshape(
              x_test.shape[0], img_rows, img_cols, 1
          )
          input_shape = (img_rows, img_cols, 1)
  
      if not is_quantized_model:
          x_train = x_train.astype("float32")
          x_test = x_test.astype("float32")
          x_train /= 255
          x_test /= 255
  
      # convert class vectors to binary class matrices
      y_train = tf.keras.utils.to_categorical(y_train, num_classes)
      y_test = tf.keras.utils.to_categorical(y_test, num_classes)
  
      return (x_train, y_train), (x_test, y_test), input_shape
  
file_path: model-optimization/tensorflow_model_optimization/python/core/keras/test_utils.py
imports: import numpy
  import tensorflow
label_desc: Use image node_id 2 for calling a global function with example usage: get_preprocessed_mnist_data(img_rows, img_cols, num_classes, is_quantized_model) and returns: x_train,  y_train,  x_test,  y_test,  input_shape