code: 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,
      )
  
file_path: trax/trax/supervised/mnist_test.py
imports: import io
  import itertools
  from unittest import mock
  from absl.testing import absltest
  from trax import layers
  from trax.data import inputs
  from trax.data import tf_inputs
  from trax.optimizers import adam
  from trax.supervised import training
label_desc: Use image node_id 3 for calling a global function with example usage: _mnist_brightness_dataset() and returns: inputs, new_stream