A pseudo-random number generator (PRNG) uses a deterministic algorithm to produce sequences that approximate random numbers. A seed initializes its internal state, and repeating that state reproduces the same sequence. PRNGs support weight initialization, data sampling, augmentation, dropout, generative models, hyperparameter tuning, Monte Carlo simulation, and statistical sampling. section 9.1 section 9.0
JAX makes the generator functional: randomized functions receive an explicit key instead of reading hidden mutable state. Use each key once, then split it into new keys for later operations. This design keeps randomness explicit and supports reproducible functional programs, unlike the usual stateful pattern associated with NumPy and commonly with TensorFlow or PyTorch. section 2.4.1 section 9.0
Initializing neural-network parameters
from jax import random key = random.PRNGKey(0) w_key, b_key = random.split(key) weights = random.normal(w_key, (512, 784)) biases = random.normal(b_key, (512,)) The separate keys make the two random draws explicit. A larger initializer can split one key into keys for every layer, as in the network-parameter example. 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
Browse the complete listing for layer sizes, parameter scaling, per-layer key splitting, and dense-layer initialization. section 2.4.1
