A key is JAX’s explicit handle for pseudorandom state. Create one with random.PRNGKey(seed), pass it to each randomized function, and split it before obtaining independent subsequent random values. Unlike NumPy’s usual hidden RNG state, a JAX random function does not update the key. section 2.4.1 section 9.2.3
Reuse versus split
key = random.PRNGKey(42) vals = random.normal(key, shape=(3, 5)) more_vals = random.normal(key, shape=(3, 5)) key1, key2 = random.split(key) vals = random.normal(key1, shape=(3, 5)) more_vals = random.normal(key2, shape=(3, 5)) The first pair is identical because both calls receive the same key. The split keys produce different random values.
Use key management to make randomness explicit:
- Split one key into several keys for separate layers, random operations, or later use; splitting into 100 keys leaves one key for further splitting and 99 subkeys for generation. section 2.4.1
Listing 2.4 Initializing the neural network
from jax import random
LAYER_SIZES = [28*28, 512, 10] #1
PARAM_SCALE = 0.01 #2
def init_network_params(sizes, key=random.PRNGKey(0), scale=1e-2):
"""Initialize all layers for a fully-connected
neural network with given sizes"""
def random_layer_params(m, n, key, scale=1e-2):
"""A helper function to randomly initialize
weights and biases of a dense layer"""
w_key, b_key = random.split(key) #3
return scale * random.normal(w_key, (n, m)),
scale * random.normal(b_key, (n,)) #4
keys = random.split(key, len(sizes))
return [random_layer_params(m, n, k, scale)
for m, n, k in zip(sizes[:-1], sizes[1:], keys)] #5
params = init_network_params(
LAYER_SIZES, random.PRNGKey(0), scale=PARAM_SCALE)
#1 The list with layer sizes
#2 Parameter for scaling random values
#3 Generates random keys (more on this in chapter 7)
#4 Generates random values for the layer parameters w and b
#5 Running generation for all the layers
Inspect key splitting in a complete neural-network parameter-initialization listing. section 2.4.1
