code: def prepare_mnist(
      batch_size: int, train_shuffle_buffer: int = 1000
  ) -> Dict[str, tf.data.Dataset]:
      """Generates a dictionary with keys measure, train, test, all from mnist.
  
      https://www.tensorflow.org/datasets/catalog/mnist
  
      Training has 60k examples, while test has 10k examples.
  
      """
      train = tfds.load(
          MNIST, split=tfds.Split.TRAIN, shuffle_files=True
      ).map(normalize)
      test = tfds.load(MNIST, split=tfds.Split.TEST).map(normalize)
  
      return {
          # This will create a buffer of `train_shuffle_buffer` items in memory,
          # then sample from these to get samples. This biases toward the first
          # batch of examples. the "repeat" effectively lets you go forever.
          #
          # If you take enough batches to run out the epoch, you'll start pulling
          # from the beginning again.
          "batched": train.shuffle(train_shuffle_buffer)
          .repeat()
          .batch(batch_size),
          # These are the unbatched train and test sets, used for evaluation and
          # metric generation at the end of each batch.
          "train": train,
          "test": test,
      }
  
imports: from __future__ import division, print_function
  import getpass
  from typing import Dict, Optional
  import fs
  import tensorflow
  import tensorflow_datasets
  import tqdm
  import uv
  import uv.tensorflow.reporter
  import uv.util
  from absl import app
  from uv.fs.reporter import FSReporter
  from uv.mlflow.reporter import MLFlowReporter
  from trainer import cli
label_desc: Use image node_id 2 for calling a global function with example usage: prepare_mnist(batch_size, train_shuffle_buffer) and returns: dict