code: def mnist(permute_train=False):
      """Download, parse and process MNIST data to unit scale and one-hot labels."""
      train_images, train_labels, test_images, test_labels = mnist_raw()
  
      train_images = _partial_flatten(train_images) / np.float32(255.0)
      test_images = _partial_flatten(test_images) / np.float32(255.0)
      train_labels = _one_hot(train_labels, 10)
      test_labels = _one_hot(test_labels, 10)
  
      if permute_train:
          perm = np.random.RandomState(0).permutation(
              train_images.shape[0]
          )
          train_images = train_images[perm]
          train_labels = train_labels[perm]
  
      return train_images, train_labels, test_images, test_labels
  
file_path: jax/examples/datasets.py
imports: import array
  import gzip
  import os
  from os import path
  import struct
  import urllib.request
  import numpy
label_desc: Use image node_id 5 for calling a global function with example usage: mnist(permute_train) and returns: train_images,  train_labels,  test_images,  test_labels