Skip to content

← Flax step 4 of 100

Medium Primitives

init() and apply() Round-Trip

Jump to the task

In Flax Linen (flax.linen), creating a module describes a computation. Calling init creates its variables, and calling apply uses those variables to compute an output. This lesson practices that workflow with a small neural network. Flax NNX has a different API.

Keep variables separate from outputs

Once you have defined your module class, the call sequence is:

model = TwoLayer(hidden=int(hidden), out=int(out))
key = jax.random.PRNGKey(int(seed))
variables = model.init(key, x)
y = model.apply(variables, x)

Here variables is a mapping containing a "params" collection. init runs the module's initialization path, creating each layer's parameters when needed. It returns the variables, not the network's output. Pass the whole mapping to apply, not just variables["params"].

For D=3, hidden=4 and out=2, the parameter shapes are:

Parameter Shape
variables["params"]["Dense_0"]["kernel"] (3, 4)
variables["params"]["Dense_0"]["bias"] (4,)
variables["params"]["Dense_1"]["kernel"] (4, 2)
variables["params"]["Dense_1"]["bias"] (2,)

The forward calculation is relu(x @ W0 + b0) @ W1 + b1. There is no activation after the second Dense layer, so outputs may be negative. You can inspect parameter shapes with jax.tree_util.tree_map(jnp.shape, variables).

Reusing initialized variables

For this network, repeated apply calls with the same variables and input produce the same output. No additional random key is needed: there is no stochastic layer such as dropout. To train later, differentiate a loss that calls apply and update the parameters; apply itself does not train them.

Dense layers use the last input dimension as the feature count. After initializing with shape (3,), the same variables also accept a batch of shape (B, 3) and return (B, out). Shape (3, 4) would fail because it has four features, not three. This exercise only asks you to return the output for the supplied one-dimensional x.

Problem

Implement two_layer_forward(seed, x, hidden, out):

  1. Define a Linen module with @nn.compact on __call__.
  2. Apply nn.Dense(features=hidden), then jax.nn.relu, then nn.Dense(features=out). Use the default Dense initializers and biases. Leave the layers unnamed so Linen assigns Dense_0 and Dense_1.
  3. Create the module and initialize it with jax.random.PRNGKey(int(seed)) and x. Do not split the key before passing it to init.
  4. Apply the initialized module to x and return its output.

x is a finite float32 JAX array of shape (D,), with D >= 1. hidden and out are positive integer-valued scalars; seed is a nonnegative integer-valued scalar. The harness may supply these three scalars as floats, so cast each to int before use. Return a JAX array of shape (out,). Do not train the model or change the input.

The initializer, layer names and key handling matter because the expected outputs use this specific initialization procedure.