Need a compact way to describe a Flax neural network? The @nn.compact annotation lets you place layer definitions and forward-pass logic together inside __call__. Flax still exposes the model functionally: initialization and application consume and return variables rather than maintaining hidden state. section 11.1.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 MLP example places two Dense layers and a Swish activation directly in __call__, then initializes variables with model.init and applies the model through its functional interface. section 11.1.1
Use @nn.compact for straightforward single-pass models. For more complex configurations, such as separate encoder and decoder methods or multiple forward passes, use the more capable setup() pattern instead. section 11.1.1
