'Args:
train_or_test (str): either \'train\' or \'test\'
shuffle (bool): shuffle the dataset'

def __init__(self, train_or_test, shuffle=True, dir=None):

    if (dir is None):
        dir = get_dataset_path('mnist_data')
    assert (train_or_test in ['train', 'test'])
    self.train_or_test = train_or_test
    self.shuffle = shuffle
    def get_images_and_labels(image_file, label_file):
        f = maybe_download(image_file, dir)
        images = extract_images(f)
        f = maybe_download(label_file, dir)
        labels = extract_labels(f)
        assert (images.shape[0] == labels.shape[0])
        return (images, labels)
    if (self.train_or_test == 'train'):
        (self.images, self.labels) = get_images_and_labels('train-images-idx3-ubyte.gz', 'train-labels-idx1-ubyte.gz')
    else:
        (self.images, self.labels) = get_images_and_labels('t10k-images-idx3-ubyte.gz', 't10k-labels-idx1-ubyte.gz')
