code: def get_data(seed):
      """
      This is the function to generate subsets of the data for training models.
  
      First, we get the training dataset either from the numpy cache
      or otherwise we load it from tensorflow datasets.
  
      Then, we compute the subset. This works in one of two ways.
  
      1. If we have a seed, then we just randomly choose examples based on
         a prng with that seed, keeping FLAGS.pkeep fraction of the data.
  
      2. Otherwise, if we have an experiment ID, then we do something fancier.
         If we run each experiment independently then even after a lot of trials
         there will still probably be some examples that were always included
         or always excluded. So instead, with experiment IDs, we guarantee that
         after FLAGS.num_experiments are done, each example is seen exactly half
         of the time in train, and half of the time not in train.
  
      """
      DATA_DIR = os.path.join(os.environ["HOME"], "TFDS")
  
      if os.path.exists(os.path.join(FLAGS.logdir, "x_train.npy")):
          inputs = np.load(os.path.join(FLAGS.logdir, "x_train.npy"))
          labels = np.load(os.path.join(FLAGS.logdir, "y_train.npy"))
      else:
          print("First time, creating dataset")
          data = tfds.as_numpy(
              tfds.load(
                  name=FLAGS.dataset, batch_size=-1, data_dir=DATA_DIR
              )
          )
          inputs = data["train"]["image"]
          labels = data["train"]["label"]
  
          inputs = (inputs / 127.5) - 1
          np.save(os.path.join(FLAGS.logdir, "x_train.npy"), inputs)
          np.save(os.path.join(FLAGS.logdir, "y_train.npy"), labels)
  
      nclass = np.max(labels) + 1
  
      np.random.seed(seed)
      if FLAGS.num_experiments is not None:
          np.random.seed(0)
          keep = np.random.uniform(
              0, 1, size=(FLAGS.num_experiments, FLAGS.dataset_size)
          )
          order = keep.argsort(0)
          keep = order < int(FLAGS.pkeep * FLAGS.num_experiments)
          keep = np.array(keep[FLAGS.expid], dtype=bool)
      else:
          keep = (
              np.random.uniform(0, 1, size=FLAGS.dataset_size)
              <= FLAGS.pkeep
          )
  
      if FLAGS.only_subset is not None:
          keep[FLAGS.only_subset :] = 0
  
      xs = inputs[keep]
      ys = labels[keep]
  
      if FLAGS.augment == "weak":
          aug = lambda x: augment(x, 4)
      elif FLAGS.augment == "mirror":
          aug = lambda x: augment(x, 0)
      elif FLAGS.augment == "none":
          aug = lambda x: augment(x, 0, mirror=False)
      else:
          raise
  
      train = DataSet.from_arrays(xs, ys, augment_fn=aug)
      test = DataSet.from_tfds(
          tfds.load(
              name=FLAGS.dataset, split="test", data_dir=DATA_DIR
          ),
          xs.shape[1:],
      )
      train = (
          train.cache()
          .shuffle(8192)
          .repeat()
          .parse()
          .augment()
          .batch(FLAGS.batch)
      )
      train = train.nchw().one_hot(nclass).prefetch(16)
      test = test.cache().parse().batch(FLAGS.batch).nchw().prefetch(16)
  
      return train, test, xs, ys, keep, nclass
  
imports: import functools
  import os
  import shutil
  from typing import Callable
  import json
  import jax
  import jax.numpy
  import numpy
  import tensorflow
  import tensorflow_datasets
  from absl import app, flags
  import objax
  from objax.jaxboard import SummaryWriter, Summary
  from objax.util import EasyDict
  from objax.zoo import convnet, wide_resnet
  from dataset import DataSet
label_desc: Use image node_id 3 for calling a global function with example usage: get_data(seed) and returns: train,  test,  xs,  ys,  keep,  nclass