Need Gaussian noise for an image augmentation? jax.random.normal() creates an array of standard normal samples: mean 0, standard deviation 1. It takes an explicit PRNG key and output shape, so the shape can match an image or any other tensor. section 9.1.2
Generate and scale image noise
seed = 42 key = jax.random.PRNGKey(seed) std_noise = jax.random.normal(key, image.shape) noise = 0.5 + 0.1 * std_noise new_image = image + noise The sampled noise is scaled and shifted before being added. Because addition can move pixel values outside [0, 1], the resulting image is normalized afterward. section 9.1.2
Keep these JAX-specific details in mind:
- Create a key with random.PRNGKey(seed); the function consumes the key explicitly. section 9.1.2
- Pass the desired output shape as the second argument; unlike NumPy, the call does not take mean and standard-deviation parameters. section 9.1.2
- Do not reuse a key when different values are required; split it and pass distinct subkeys to successive calls. section 9.2.3
Listing 9.5 Generating the random tensor
seed = 42 #1 key = jax.random.PRNGKey(seed) #2 std_noise = jax.random.normal(key, image.shape) #3 std_noise.min(), std_noise.max() #4 >>> (Array(-4.170905, dtype=float32), Array(4.259979, dtype=float32)) noise = 0.5 + 0.1*std_noise #5 noise.min(), noise.max() #6 >>> (Array(0.08290949, dtype=float32), Array(0.92599785, dtype=float32)) plt.imshow(noise)
Shows the complete key setup, normal sampling, observed range, and affine scaling used to prepare image noise. section 9.1.2
