nn.Sequential Composition

nn.Sequential applies a list of callables in order. Use it here to build a three-layer multilayer perceptron (MLP): two hidden Dense layers with ReLU activations, followed by a Dense output layer.

A chain of layers and functions

class SeqMLP(nn.Module):
    hidden: int
    out: int

    @nn.compact
    def __call__(self, x):
        return nn.Sequential([
            nn.Dense(features=self.hidden),
            jax.nn.relu,
            nn.Dense(features=self.hidden),
            jax.nn.relu,
            nn.Dense(features=self.out),
        ])(x)

Dense instances own learned parameters. jax.nn.relu is a function with no parameters; it replaces negative values with zero. Both can be entries in the same Sequential list. Pass the function itself, not jax.nn.relu().

For a vector input of length D, the intermediate shapes are (hidden,), (hidden,) and (out,) after the three Dense layers. ReLU preserves shape. There is no final activation, so the output may contain negative values.

Parameter paths in this construction

The Dense children above are constructed inside the enclosing compact method before the Sequential object is created. They belong to that enclosing module. With hidden=4, out=2 and input length 3, the parameters are:

params
  Dense_0: kernel (3, 4), bias (4,)
  Dense_1: kernel (4, 4), bias (4,)
  Dense_2: kernel (4, 2), bias (2,)

The two ReLU functions add no parameter entries. This particular model has no Sequential_0/layers_0 parameter subtree. Constructing a standalone Sequential, or creating children inside other scopes, can give different paths. Inspect the initialized mapping rather than assuming one naming scheme applies to every construction.

An explicit compact method that creates the same three unnamed Dense layers in order and inserts ReLU between them has matching parameter paths. For this pair of definitions, the same seed gives the same parameters and output. Changing parameter paths can change seeded initialization.

How values move through the chain

For this MLP, each item receives and returns one array. Sequential also supports structured intermediate results: a tuple is unpacked as positional arguments to the next callable, and a dictionary as keyword arguments. For example, a function returning (x, x) can feed a function accepting two arrays and adding them.

Initial arguments and keyword arguments go to the first callable. They are not automatically forwarded to every later layer. Adding a train=False argument to apply does not, by itself, configure an arbitrary BatchNorm later in a chain. Set its behavior when constructing it or use a wrapper that passes the flag where needed.

Explicit forward code is often easier to follow when values branch, skip layers or need to be saved for later. Sequential is convenient when the output of each step naturally supplies the next step's inputs.