code: def build_train_set(batch_size, ds_builder):
      """Builds train dataset."""
  
      train_ds = ds_builder.as_dataset(split=tfds.Split.TRAIN)
      train_ds = train_ds.map(prepare_image)
      train_ds = train_ds.cache()
      train_ds = train_ds.repeat()
      train_ds = train_ds.shuffle(50000)
      train_ds = train_ds.batch(batch_size)
      train_ds = iter(tfds.as_numpy(train_ds))
      return train_ds
  
imports: import jax
  import jax.numpy
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 1 for calling a global function with example usage: build_train_set(batch_size, ds_builder) and returns: train_ds