Random augmentation chooses one transformation for an image at runtime, such as adding noise or flipping horizontally. In JAX, jax.lax.switch selects the branch while separate PRNG keys preserve independent randomness. section 9.1.3
A practical JAX implementation follows these rules:
- Give every augmentation the same signature so lax.switch can call them uniformly. section 9.1.3
- Split the input key before choosing an augmentation and using randomness inside that branch; reusing one key entangles the results. section 9.1.3
- Batch the operation with vmap, passing images and keys along axis zero while keeping the augmentation list unbatched. section 6.3.1
Noise or horizontal flip
Define add_noise_func and horizontal_flip_func, place them in augmentations, then let random_augmentation split rng_key, draw an index with jax.random.randint, and dispatch through jax.lax.switch. For a batch, split batch_size + 1 keys, retain one updated key, and pass the subkeys to vmap. section 9.1.3 section 9.3.1
Listing 9.22 Image augmentation loop
key = random.PRNGKey(42) #1
for x, y in train_data:
display_batch(x, y, 4, 8, CLASS_NAMES) #2
batch_size = len(x)
key, *subkeys = random.split(key, num=batch_size+1) #3
aug_x = jax.vmap(
random_augmentation, in_axes=(0,None,0) #4
)(x, augmentations, jnp.array(subkeys)) #5
display_batch(aug_x, y, 4, 8, CLASS_NAMES) #6
# ...
# do some other processing and neural network training
# ...
The batch-loop listing makes key management and per-image vectorization concrete, including the updated key needed for the next batch. section 9.3.1
