code: def build_test_set(ds_builder):
      """Builds train dataset."""
      test_ds = ds_builder.as_dataset(split=tfds.Split.TEST)
      test_ds = test_ds.map(prepare_image).batch(10000)
      test_ds = jnp.array(list(test_ds)[0])
      test_ds = jax.device_put(test_ds)
      return test_ds
  
imports: import jax
  import jax.numpy
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 2 for calling a global function with example usage: build_test_set(ds_builder) and returns: test_ds