Use lax.switch when an integer index must select one function from several alternatives. Unlike ordinary Python branching, it provides JAX control flow that is JIT-able and differentiable. section 3.4.1
Random image augmentation
Choose a random index with random.randint, then call lax.switch(augmentation_index, augmentations, image). The selected function can add noise, flip, rotate, or adjust colors. For batches, split the random key per image and combine the function with vmap; reusing one key repeats the same transformation. section 3.4.1 section 6.3.1
Listing 3.27 Control flow example with lax.switch
from jax import random #1
def random_augmentation(image, augmentations, rng_key):
'''A function that applies a random transformation to an image'''
augmentation_index = random.randint( #2
key=rng_key, minval=0, #2
maxval=len(augmentations), shape=()) #2
augmented_image = lax.switch(
augmentation_index, augmentations, image) #3
return augmented_image
new_image = random_augmentation(image, augmentations, random.PRNGKey(42))
#1 Imports module for random numbers in JAX (more on this in chapter 9)
#2 Generates a random integer value to index the augmentations list
#3 Uses lax.switch() function to choose between multiple options
Runnable pattern for selecting a random augmentation with lax.switch. section 3.4.1
Listing 6.19 Augmenting multiple elements of data
images = jnp.repeat( jnp.reshape(image, (1,len(image))), 10, axis=0) #1 images.shape >>> (10, 100) rng_keys = random.split(random.PRNGKey(211), num=len(images)) #2 random_augmentation_batch = jax.vmap( random_augmentation, in_axes=(0,None,0)) #3 augmented_images = random_augmentation_batch( images, augmentations, rng_keys) #4 #1 Generates a batch of stub data #2 Generates an array of random number generator keys #3 Auto-vectorizes the function #4 Applies the vectorized function to a batch
Batch pattern showing per-element keys, in_axes, and vectorized augmentation. section 6.3.1
