code: def get_dataset():
      (x_train, y_train), (
          x_test,
          y_test,
      ) = tf.keras.datasets.mnist.load_data()
      x_train = (x_train / 255).astype("float32")
      x_test = (x_test / 255).astype("float32")
      return (x_train, y_train), (x_test, y_test)
  
imports: import os
  import shutil
  from absl.testing import parameterized
  import tensorflow
  from tensorflow_model_optimization.python.core.common.keras.compression.algorithms import epr
label_desc: Use image node_id 2 for calling a global function with example usage: get_dataset() and returns: x_train,  y_train,  x_test,  y_test