code: 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)
  
file_path: trax/trax/tf_numpy/examples/mnist/dataset.py
imports: import gzip
  import os
  import pickle
  import random
  import urllib
  import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: load() and returns: shuffle,  validation_data,  test_data, zip