Skip to content

← Flax NNX step 21 of 100

Medium Primitives

NNX Implement Dense

Jump to the task

Why this matters

A dense layer computes an affine transformation: y = x @ W + b. Writing one as an NNX module shows how a constructor creates parameters and a forward method uses them. This exercise uses a specific random initialization so its outputs are reproducible for a given seed.

Shapes and parameters

For an input vector of length D_in, the kernel has shape (D_in, D_out) and the bias has shape (D_out,). Each output is a weighted sum of the input entries plus one bias value. There is no activation function in this layer.

Wrap both arrays in nnx.Param to identify them as parameters for filters and the usual NNX differentiation/optimization workflow. A raw array can still appear in NNX state; it simply does not match an nnx.Param filter.

The initialization used here

Draw one standard normal array and multiply it by 1 / sqrt(D_in). Each weight then has theoretical variance 1 / D_in. Under suitable independence and zero-mean assumptions, this scale helps keep activation variance from growing solely because a layer has more inputs. It is not a guarantee of stable activations in every network.

This is an untruncated normal draw. Do not substitute jax.nn.initializers.lecun_normal() or the default nnx.Linear initializer: the built-in LeCun normal initializer uses a scaled truncated normal, so its seeded values differ from the recipe required here.

A complete example

import jax
import jax.numpy as jnp
from flax import nnx

class MyDense(nnx.Module):
    def __init__(self, in_features, out_features, rngs):
        key = rngs.params()
        kernel = jax.random.normal(
            key, (in_features, out_features), dtype=jnp.float32
        ) * (1.0 / jnp.sqrt(in_features))
        self.kernel = nnx.Param(kernel)
        self.bias = nnx.Param(jnp.zeros(out_features, dtype=jnp.float32))

    def __call__(self, x):
        return x @ self.kernel[...] + self.bias[...]

model = MyDense(3, 4, nnx.Rngs(0))
x = jnp.array([1.0, 2.0, 3.0], dtype=jnp.float32)
assert model(x).shape == (4,)
assert model.kernel[...].shape == (3, 4)
assert model.bias[...].shape == (4,)
assert len(jax.tree.leaves(nnx.state(model, nnx.Param))) == 2

# Set known parameters to check the affine calculation.
model.kernel[...] = jnp.ones((3, 4), dtype=jnp.float32)
model.bias[...] = jnp.array([0.0, 1.0, 2.0, 3.0], dtype=jnp.float32)
assert jnp.allclose(model(x), jnp.array([6.0, 7.0, 8.0, 9.0]))

The final assignments only demonstrate the calculation. Keep the specified random kernel and zero bias when solving the exercise. Request exactly one rngs.params() key for the kernel; zero bias initialization needs no key.

Problem

Implement dense_forward(seed, x, features):

  1. Define MyDense with the two Param attributes and forward method above.
  2. Construct nnx.Rngs(int(seed)) and initialize MyDense using input width x.shape[-1] and output width int(features).
  3. Return the model's output on x.

Use float32 for the kernel and bias. Reproduce the untruncated normal draw and scaling operation rather than calling nnx.Linear. If JIT-compiling the entire entrypoint, treat seed and features as static constructor arguments.

Inputs: seed is a nonnegative integer-valued scalar and features is a positive integer-valued scalar, both supplied as floats and cast to int. x is a nonempty finite float32 vector. Values are small enough for a finite forward result.

Output: vector of shape (int(features),).