Initializes a Module’s variables from dummy inputs and a PRNG key. The returned variables include parameters, and for stateful models can include model state too. section 11.1.1 section 11.4
It usually appears right after you instantiate a Module and before any forward pass. Flax uses the dummy input to infer shapes, so you declare features rather than manual weight shapes. section 11.1.1 section 11.4
Typical uses here are:
- create the initial parameter pytree for a model. section 11.1.1
- include non-parameter state such as BatchNorm statistics when present. section 11.2.1
- inspect the inferred structure with tabulate() after initialization. section 11.2.1
Minimal MLP initialization
An MLP instance is created, a random flattened image is generated, and model.init(key, random_flattened_image) returns the variables. The same pattern can be paired with init_with_output() when the forward-pass output is also needed. section 11.1.1
Listing 11.1 Initializing and applying MLP in pure JAX (reproducing listings 2.4 and 2.5)
from jax import random
import jax.numpy as jnp
from jax.nn import swish
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)
return (scale * random.normal(w_key, (n, m)),
scale * random.normal(b_key, (n,))) #3
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)] #4
params = init_network_params(
LAYER_SIZES, random.PRNGKey(0), scale=PARAM_SCALE)
def predict(params, image):
"""Function for per-example predictions."""
activations = image #5
for w, b in params[:-1]: #6
outputs = jnp.dot(w, activations) + b #7
activations = swish(outputs) #7
final_w, final_b = params[-1] #8
logits = jnp.dot(final_w, activations) + final_b #8
return logits
Shows the pure-JAX equivalent that init() replaces by inferring parameters from explicit layer sizes and random keys. section 11.1.1
