Need several random draws without hidden mutable state? In JAX, split divides one explicit PRNG key into separate keys, so each randomized operation receives its own input. section 9.2.3 section 2.4.1
Call jax.random.split(key, num), which deterministically returns the requested number of new keys; the default number is two. Treat the original key as consumed and use each resulting key once. section 9.2.3 section 9.4
Use splitting when random operations need independent key inputs:
- Split separate keys for separate random draws, such as generating two random tensors. section 9.2.3
- Split one key for augmentation selection and another for noise generation; reusing one key can entangle the results. section 9.1.3
- Split by layer, then split each layer key again for weights and biases during neural-network initialization. section 9.3.2
Two independent draws
key1, key2 = random.split(key, num=2) vals = random.normal(key1, shape=(3, 5)) more_vals = random.normal(key2, shape=(3, 5)) Passing the same key twice would reproduce the same values; splitting supplies distinct keys for distinct draws. section 9.2.3
Listing 9.23 Neural network random initialization
import jax.numpy as jnp
from jax import random
LAYER_SIZES = [200*200*3, 2048, 1024, 2] #1
PARAM_SCALE = 0.01 #2
def random_layer_params(m, n, key, scale=1e-2): #3
w_key, b_key = random.split(key) #4
return (scale * random.normal(w_key, (n, m)),
scale * random.normal(b_key, (n,))) #5
def init_network_params(sizes,
key=random.PRNGKey(0), scale=0.01): #6
keys = random.split(key, len(sizes)-1) #7
return [random_layer_params(m, n, k, scale)
for m, n, k in zip(sizes[:-1], sizes[1:], keys)] #8
key = random.PRNGKey(42) #9
params = init_network_params(LAYER_SIZES, key, scale=PARAM_SCALE) #10
#1 Describes the structure of a neural network
#2 Standard deviation for neural network weights
#3 A function to initialize weights and biases of a dense layer
#4 Splits a key into a key for weights and a key for biases
#5 Consumes both keys to generate weights and biases
#6 A function to initialize all layers for a fully connected neural network with given sizes
#7 Splits a key into many keys equal to the number of layers
#8 Passes each key to the function for layer initialization
#9 Generates an initial key with some seed
#10 Initializes all the layers
Inspect a complete initialization pattern that splits a network key across layers and then splits each layer key for weights and biases. section 9.3.2
