Module with setup() (alternative to compact)

Flax Linen lets you define a child layer in setup, then call it from a forward method. This separates the list of layers from the calculation that uses them. In this exercise, the model contains one Dense layer.

Define the child, then call it

class DenseSetup(nn.Module):
    features: int

    def setup(self):
        self.dense = nn.Dense(features=self.features)

    def __call__(self, x):
        return self.dense(x)

features is a configuration field. self.dense is the child module. Its forward calculation is x @ kernel + bias, with no activation. For input shape (fan_in,), the kernel has shape (fan_in, features) and the output has shape (features,).

Flax invokes setup lazily when it needs the bound module, such as during init or apply. Do not call setup yourself. init creates the parameter arrays; apply uses the arrays you pass to it. Calling setup as part of that process does not mean the parameters are randomly initialized again.

Names affect initialization

Assigning the unnamed Dense to self.dense gives it the child name dense. Its arrays appear at:

variables["params"]["dense"]["kernel"]
variables["params"]["dense"]["bias"]

An unnamed Dense constructed inside @nn.compact usually has a generated name such as Dense_0. These paths affect the random keys Flax derives for initialization. Two models can have the same architecture and input seed but different initial weights because their child names differ.

To express this particular model in compact style with matching names, use:

class DenseCompact(nn.Module):
    features: int

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

These two definitions produce the same parameters and outputs for the same seed and input. The explicit name matters here, not which style you choose.

Choosing a style

setup is useful when several forward methods need access to a child through self. Compact keeps a layer's definition next to the calculation that uses it and can use input shapes when defining parameters. Both styles support explicit names and reuse of an existing layer instance.

You can also use setup and a compact method in the same class. For example, define a reusable encoder in setup, then create an output layer inside a compact forward method. Give distinct layers distinct names. This exercise uses only setup so you can practice that pattern on its own.

Neither style exempts Python code from JAX tracing rules. Branching on fixed configuration, such as the number of layers, is different from branching on a runtime array value. When a function is jitted, its Python setup logic can run while JAX traces the function; do not rely on it running on every invocation of the compiled function.