code: def load_mnist_dataset():
      mnist = keras.datasets.mnist
      (train_images, train_labels), (
          test_images,
          test_labels,
      ) = mnist.load_data()
  
      # Normalize the input image so that each pixel value is between 0 to 1.
      train_images = train_images / 255.0
      test_images = test_images / 255.0
  
      return (train_images, train_labels), (test_images, test_labels)
  
file_path: model-optimization/tensorflow_model_optimization/python/examples/clustering/keras/mnist/mnist_cnn.py
imports: from __future__ import print_function
  import datetime
  import os
  from absl import app
  from absl import flags
  import tensorflow
  from tensorflow_model_optimization.python.core.clustering.keras import cluster
  from tensorflow_model_optimization.python.core.clustering.keras import cluster_config
  from tensorflow_model_optimization.python.core.clustering.keras import clustering_callbacks
label_desc: Use image node_id 1 for calling a global function with example usage: load_mnist_dataset() and returns: train_images,  train_labels,  test_images,  test_labels