def load():
  """Loads the dataset.

  Looks for the dataset at /tmp/mnist.pkl.gz and downloads it if it is not there
  already.

  Note: The training data is shuffled.

  Returns:
    ((train_x, train_y), (valid_x, valid_y), (test_x, test_y)).
    Shapes:
      train_x: num_training_examples x image_size
      train_y: num_training_examples x num_classes
      valid_x: num_validation_examples x image_size
      valid_y: num_validation_examples x num_classes
      test_x: num_test_examples x image_size
      test_y: num_test_examples x num_classes
  """
  filepath = _maybe_download()
  with gzip.open(os.path.join(filepath), 'rb') as f:
    training_data, validation_data, test_data = pickle.load(f)
  training_data = (training_data[0], [to_one_hot(x) for x in training_data[1]])
  validation_data = (validation_data[0],
                     [to_one_hot(x) for x in validation_data[1]])
  test_data = (test_data[0], [to_one_hot(x) for x in test_data[1]])

  def shuffle(data):
    zipped = zip(*data)
    random.shuffle(zipped)
    return zip(*zipped)

  return (shuffle(training_data), validation_data, test_data)
Loads the dataset.

  Looks for the dataset at /tmp/mnist.pkl.gz and downloads it if it is not there
  already.

  Note: The training data is shuffled.

  Returns:
    ((train_x, train_y), (valid_x, valid_y), (test_x, test_y)).
    Shapes:
      train_x: num_training_examples x image_size
      train_y: num_training_examples x num_classes
      valid_x: num_validation_examples x image_size
      valid_y: num_validation_examples x num_classes
      test_x: num_test_examples x image_size
      test_y: num_test_examples x num_classes
  
trax/tf_numpy/examples/mnist/dataset.py