code: def mnist(split: str, batch_size: int) -> tf.data.Dataset:
      """Returns a tf.data.Dataset with MNIST image/label pairs."""
  
      def preprocess_dataset(images, labels):
          # Mnist images are int8 [0, 255], we cast and rescale to float32 [-1, 1].
          images = ((tf.cast(images, tf.float32) / 255.0) - 0.5) * 2.0
          return images, labels
  
      dataset = tfds.load(
          name="mnist",
          split=split,
          shuffle_files=split == "train",
          as_supervised=True,
      )
      dataset = dataset.map(preprocess_dataset)
      dataset = dataset.shuffle(buffer_size=4 * batch_size)
      dataset = dataset.batch(batch_size)
      # Cache the result of the data pipeline to avoid recomputation. The pipeline
      # is only ~100MB so this should not be a significant cost and will afford a
      # decent speedup.
      dataset = dataset.cache()
      # Prefetching batches onto the GPU will help avoid us being too input bound.
      # We allow tf.data to determine how much to prefetch since this will vary
      # between GPUs.
      dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
      return dataset
  
file_path: sonnet/examples/simple_mnist.py
imports: from typing import Dict
  from absl import app
  import sonnet
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 1 for calling a global function with example usage: mnist(split, batch_size) and returns: dataset, images,  labels