Module with @nn.compact

Flax Linen is a library for defining neural networks in JAX. This exercise introduces a model with one Dense layer and the three steps needed to use it: construct the model, initialize its variables, then apply it to an input.

Define a module

class DenseCompact(nn.Module):
    features: int

    @nn.compact
    def __call__(self, x):
        return nn.Dense(features=self.features)(x)

features is a configuration field: it chooses the number of output units. It is not a learned parameter. The @nn.compact decorator lets this method create a child Dense layer where it is used. nn.Dense(...)(x) first constructs the child, then calls it on x.

For a vector x of shape (fan_in,), Dense computes x @ kernel + bias. The kernel has shape (fan_in, features), the bias has shape (features,), and the result has shape (features,). No activation is included here, so negative output values are allowed.

Construct, initialize, apply

import jax
import jax.numpy as jnp
import flax.linen as nn

model = DenseCompact(features=8)
x = jnp.ones((16,), dtype=jnp.float32)
key = jax.random.PRNGKey(0)
variables = model.init(key, x)
y = model.apply(variables, x)

Constructing model records its configuration. init runs the model's initialization path and returns a mapping of variable collections. This model has only the params collection, which contains the learned arrays:

params
  Dense_0
    kernel: shape (16, 8)
    bias:   shape (8,)

Pass the whole variables mapping to apply, including its outer params key. apply uses those arrays to compute the output; it does not train the model or draw new weights. Here y has shape (8,).

Other models can have additional collections, such as running statistics. They may request updated collections through apply(..., mutable=...). This Dense-only exercise needs neither mutable state nor an optimizer.

Parameter names and input shapes

The first unnamed Dense child in this compact method is called Dense_0. Assigning it to a Python local variable such as layer does not change that name. A second unnamed Dense instance would be Dense_1; reusing the same instance shares its parameters instead. An explicit name= can override automatic naming, but leave it unset in this exercise because parameter paths affect seeded initialization.

The input's feature width determines the kernel's first dimension. A dummy vector with the wrong width can initialize successfully, but it creates a kernel for that width. Applying those variables to a different feature width then fails. Use the supplied x for both initialization and application.

In the usual init/apply workflow, variables are passed separately from the module. Linen also offers bound modules that hold access to variables for interactive use. Keeping arrays explicit makes it easier to use JAX transformations with functions such as model.apply; it does not remove their shape or tracing requirements.