Want a neural network that reads like a familiar Python model without giving up JAX transformations? Linen, or flax.linen, provides Flax’s recommended Module API for defining networks in a functional JAX style. section 11.1.1 section 11.0
A Linen Module groups a model’s structure and parameters in one place. You define layers in a class, commonly with @nn.compact and __call__, then use model.init to create variables and model.apply to run the network. The class feels stateful and object-oriented, while the surrounding interface remains pure: variables are passed in and returned rather than hidden in side effects. section 11.1.1
A typical Linen workflow is:
- Define the architecture as an nn.Module, placing Dense, convolutional, activation, or other standard layers in its call path. section 11.1.1 section 11.0
- Initialize variables with a JAX random key and representative input, then inspect their tree structure or tabulated model summary. section 11.1.1
- Apply the model inside loss and gradient functions, using JAX transformations such as jit and value_and_grad. section 11.1.1
- Pass non-parameter state explicitly; BatchNorm, for example, keeps running statistics in a batch_stats collection and updates them during training. section 11.2.1
Listing 11.2 Defining and applying MLP in Flax
from jax import random #1
from flax import linen as nn #1
class MLP(nn.Module): #2
"""A simple MLP model."""
@nn.compact #3
def __call__(self, x):
x = nn.Dense(features=512)(x) #4
x = nn.activation.swish(x) #4
x = nn.Dense(features=10)(x) #4
return x #4
model = MLP() #5
key1, key2 = random.split(random.PRNGKey(0))
random_flattened_image = random.normal(key1, (28*28*1,)) #6
params = model.init(key2, random_flattened_image) #7
jax.tree_util.tree_map(lambda x: x.shape, params) #8
>>> dict({
>>> params: {
>>> Dense_0: {
>>> bias: (512,),
>>> kernel: (784, 512),
>>> },
>>> Dense_1: {
>>> bias: (10,),
>>> kernel: (512, 10),
>>> },
>>> },
>>> })
model.apply(params, random_flattened_image) #9
>>> Array([ 0.43500084, 0.20896661, -0.9779186 , 0.3447541 , 0.05475511,
>>> 0.29811084, -0.24472891, -0.30921277, 0.2945502 , -0.34386003],
>>> dtype=float32)
#1 Necessary imports
#2 Class for our MLP, subclass of flax.linen.Module
#3 Uses nn.compact declaration mode
#4 Defines a net directly within a single forward-pass method
#5 Instantiates an object for the MLP
#6 Prepares a dummy input (you may also use an array of zeros as well)
#7 Initializes the model parameters
#8 Checks output shapes for the model parameters
#9 Applies the model to input data
The compact MLP listing shows the essential Linen pattern: declare a Module, compose Dense layers, initialize variables, and inspect parameter shapes. section 11.1.1
From MLPs to stateful vision models
An MNIST-style MLP uses Dense layers and swish activation, replacing scattered pure-JAX initialization and prediction code with a single model definition. The same style scales to an 18-layer ResNet, where convolution, normalization, residual blocks, parameter collections, and model state feed a JIT-compiled training step. section 11.1.1 section 11.2.1
