code: def get_dataset(
      name,
      n_train=None,
      n_test=None,
      permute_train=False,
      do_flatten_and_normalize=True,
      data_dir=None,
      input_key="image",
  ):
      """Download, parse and process a dataset to unit scale and one-hot labels."""
      # Need this following http://cl/378185881 to prevent GPU test breakages.
      tf.config.set_visible_devices([], "GPU")
  
      ds_builder = tfds.builder(name)
  
      ds_train, ds_test = tfds.as_numpy(
          tfds.load(
              name + (":3.*.*" if name != "imdb_reviews" else ""),
              split=[
                  "train"
                  + ("[:%d]" % n_train if n_train is not None else ""),
                  "test"
                  + ("[:%d]" % n_test if n_test is not None else ""),
              ],
              batch_size=-1,
              as_dataset_kwargs={"shuffle_files": False},
              data_dir=data_dir,
          )
      )
  
      train_images, train_labels, test_images, test_labels = (
          ds_train[input_key],
          ds_train["label"],
          ds_test[input_key],
          ds_test["label"],
      )
  
      if do_flatten_and_normalize:
          train_images = _partial_flatten_and_normalize(train_images)
          test_images = _partial_flatten_and_normalize(test_images)
  
      num_classes = ds_builder.info.features["label"].num_classes
      train_labels = _one_hot(train_labels, num_classes)
      test_labels = _one_hot(test_labels, num_classes)
  
      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: neural-tangents/examples/datasets.py
imports: import gzip
  import os
  import shutil
  import urllib.request
  from jax import random
  import numpy
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 3 for calling a global function with example usage: get_dataset(name, n_train, n_test, permute_train, do_flatten_and_normalize, data_dir, input_key) and returns: train_images,  train_labels,  test_images,  test_labels