code: def prepare_datasets(hparams):
      """Downloads the MNIST dataset and converts to `tf.data.Dataset` format."""
      (x_train, y_train), (
          x_test,
          y_test,
      ) = tf.keras.datasets.mnist.load_data()
  
      def make_dataset(x, y, shuffle=False):
          x = x.reshape((-1, 28, 28, 1)).astype("float32") / 255.0
          dataset = tf.data.Dataset.from_tensor_slices((x, y))
          if shuffle:
              dataset = dataset.shuffle(hparams.buffer_size)
          return dataset.batch(hparams.batch_size)
  
      return make_dataset(x_train, y_train, True), make_dataset(
          x_test, y_test
      )
  
file_path: neural-structured-learning/neural_structured_learning/examples/adv_keras_cnn_mnist.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl import app
  from absl import flags
  import attr
  import neural_structured_learning
  import tensorflow
label_desc: Use image node_id 2 for calling a global function with example usage: prepare_datasets(hparams) and returns: make_dataset,  make_dataset, dataset