def _mnist_dataset():
  """Loads (and caches) the standard MNIST data set."""
  streams = tf_inputs.data_streams('mnist')
  return inputs.batcher(streams, variable_shapes=False,
                        batch_size_per_device=256,
                        eval_batch_size=256)
Loads (and caches) the standard MNIST data set.
trax/supervised/mnist_test.py