def _mnist_brightness_dataset():
  """Loads (and caches) a MNIST mean brightness data set."""
  def preprocess_stream(stream):
    def new_stream():
      for (image, _) in stream():
        yield (image, (image / 255).mean()[None])
    return new_stream

  streams = tuple(map(preprocess_stream, tf_inputs.data_streams('mnist')))
  return inputs.batcher(streams, variable_shapes=False,
                        batch_size_per_device=256,
                        eval_batch_size=256)
Loads (and caches) a MNIST mean brightness data set.
trax/supervised/mnist_test.py