code: def load_mnist():
      partial_flatten = lambda x: np.reshape(
          x, (x.shape[0], np.prod(x.shape[1:]))
      )
      one_hot = lambda x, k: np.array(
          x[:, None] == np.arange(k)[None, :], dtype=int
      )
      (
          train_images,
          train_labels,
          test_images,
          test_labels,
      ) = data_mnist.mnist()
      train_images = partial_flatten(train_images) / 255.0
      test_images = partial_flatten(test_images) / 255.0
      train_labels = one_hot(train_labels, 10)
      test_labels = one_hot(test_labels, 10)
      N_data = train_images.shape[0]
  
      return (
          N_data,
          train_images,
          train_labels,
          test_images,
          test_labels,
      )
  
file_path: autograd/examples/data.py
imports: from __future__ import absolute_import
  import matplotlib.pyplot
  import matplotlib.image
  import autograd.numpy
  import autograd.numpy.random
  import data_mnist
label_desc: Use image node_id 1 for calling a global function with example usage: load_mnist() and returns: N_data,  train_images,  train_labels,  test_images,  test_labels