Building an MLP with flax.linen.Dense

Building an MLP with flax.linen.Dense | CodeSignal Learn

Introduction

Welcome back to the second lesson of Beyond Pure JAX: Flax & Optax for Elegant ML! You've made excellent progress so far. Having mastered the foundational concepts of Flax modules with setup() and__call__() methods in our previous lesson, you're now ready to explore one of Flax's most essential building blocks: the Dense layer . While creating custom modules gives us a deep understanding and flexibility, real-world machine learning relies heavily on proven, optimized components. The flax.linen.Dense layer is precisely such a component — a fully connected layer that handles weight initialization, bias terms, and efficient linear transformations automatically. Today, we'll construct a complete Multi-Layer Perceptron (MLP) using these built-in Dense layers, learning how to compose them elegantly while exploring two distinct approaches to module definition. By the end of this lesson, you'll confidently build neural networks using Flax's high-level abstractions, understanding both the setup() method and the powerful @nn.compact decorator. We'll apply these techniques to construct an MLP capable of learning the XOR function, demonstrating how Flax simplifies complex model construction while maintaining JAX's functional programming principles.

Understanding Flax's Built-in Dense Layer

y = Wx + b

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal

The flax.linen.Dense layer represents a fundamental building block in neural networks, implementing the ubiquitous linear transformation:

y = Wx + b

where $W$ is the weight matrix, $x$ is the input vector, and $b$ is the bias term. Unlike our custom implementations from the previous lesson, Dense layers come with sophisticated initialization strategies, automatic shape inference, and optimized computation paths.

What makes Flax's Dense layer particularly elegant is its simplicity of use combined with powerful underlying functionality. When we specify nn.Dense(features=3), we're creating a layer that will automatically determine input dimensions from the first forward pass, initialize weights using proven strategies like Xavier/Glorot initialization, and handle the linear transformation efficiently using JAX's optimized operations.

The features parameter determines the output dimensionality, that is, how many neurons this layer contains. For our XOR problem, we might use Dense(3) for a hidden layer with three neurons and Dense(1) for the output layer producing a single prediction. Flax handles all the complexity of weight matrices, bias vectors, and shape compatibility automatically, allowing us to focus on architecture design rather than implementation details.

Building MLPs with the setup() Method

Let's construct our first MLP using the familiar setup() method pattern you learned previously. This approach explicitly defines all layers during module initialization, providing clear separation between structure definition and forward pass logic. Here's how we define an MLP class using this methodology: Python import flax.linen as nn class MLP(nn.Module): hidden_features: int # Number of neurons in the hidden layer output_features: int # Number of neurons in the output layer def setup(self): # Define the layers self.dense_hidden = nn.Dense(features=self.hidden_features, name="HiddenLayer") self.dense_output = nn.Dense(features=self.output_features, name="OutputLayer") In this setup() method, we create two Dense layers as instance attributes. The hidden_features and output_features are dataclass-style attributes that parameterize our module, making it reusable for different architectures. Notice how we explicitly name our layers using the name parameter — this creates meaningful parameter names in our model's parameter tree, making debugging and analysis much easier.

Defining the Forward Pass with call

The forward pass logic comes in the__call__method, where we chain our layers together with activation functions: Python def __call__(self, x): # Define the forward pass x = self.dense_hidden(x) x = nn.sigmoid(x) # Activation for the hidden layer x = self.dense_output(x) x = nn.sigmoid(x) # Activation for the output layer return x This implementation demonstrates the power of composition in Flax. Each Dense layer call self.dense_hidden(x) performs the linear transformation, while nn.sigmoid(x) applies the activation function. The flow is intuitive: input \to hidden layer \to activation \to output layer \to final activation. The sequential application of transformations mirrors the mathematical formulation of neural networks, making the code both readable and maintainable.

The @nn.compact Alternative

Flax offers a more concise approach through the @nn.compact decorator, which allows us to define layers inline within the__call__method. This approach can be more readable for certain architectures and enables dynamic layer creation based on input properties: Python class MLPCompact(nn.Module): hidden_features: int output_features: int @nn.compact def __call__(self, x): # Define layers inline as they are used x = nn.Dense(features=self.hidden_features, name="HiddenLayer")(x) x = nn.sigmoid(x) x = nn.Dense(features=self.output_features, name="OutputLayer")(x) x = nn.sigmoid(x) return x The @nn.compact decorator transforms our__call__method into both a layer definer and a forward pass executor. When we write nn.Dense(features=self.hidden_features, name="HiddenLayer")(x), we're simultaneously creating the Dense layer and applying it to input x . This pattern is particularly useful for models where the architecture depends on input shapes or for rapid prototyping. Both approaches produce identical functionality — the choice between setup() and @nn.compact often comes down to personal preference and specific use cases. The setup() method provides explicit structure definition, which is useful for complex architectures with shared layers, while @nn.compact offers more concise, mathematical notation-like code that's perfect for straightforward sequential models.

Initializing Parameters and Understanding Structure

Now let's instantiate our models and examine how Flax handles parameter initialization. We'll use again the XOR dataset to demonstrate the process, creating both MLP variants and initializing them with identical random keys to ensure reproducible comparisons: Pythonimport jax import jax.numpy as jnp # Prepare XOR input data xor_X = jnp.array([[0., 0.], [0., 1.], [1., 0.], [1., 1.]], dtype=jnp.float32) # Initialize both MLP models xor_mlp_setup = MLP(hidden_features=3, output_features=1) xor_mlp_compact = MLPCompact(hidden_features=3, output_features=1) # Initialize parameters for both models with the same key key = jax.random.key(42) variables_setup = xor_mlp_setup.init(key, xor_X) params_setup = variables_setup['params'] # Reuse the same key to demonstrate deterministic behaviour variables_compact = xor_mlp_compact.init(key, xor_X) params_compact = variables_compact['params'] The init() method performs the crucial task of parameter initialization. It takes a PRNG key and sample input data, uses the input to infer shapes, and creates the initial parameter values. Let's examine what these initialized parameters look like: text{'HiddenLayer': {'kernel': Array([[-1.2377772 , -0.19170533, -0.26887625], [-0.6895247 , 0.5331821 , -0.06674385]], dtype=float32), 'bias': Array([0., 0., 0.], dtype=float32)}, 'OutputLayer': {'kernel': Array([[-0.2270762 ], [-0.11355998], [-0.24731739]], dtype=float32), 'bias': Array([0.], dtype=float32)}} The parameter structure reveals Flax's automatic organization. Each Dense layer contributes a kernel (weight matrix) and bias vector. The HiddenLayer has a kernel with shape (2, 3) — transforming 2 input features to 3 hidden units — and a bias of shape (3,). Similarly, the OutputLayer transforms 3 hidden units to 1 output with shapes (3, 1) and (1,), respectively. This hierarchical organization makes it easy to inspect, modify, or save specific parts of the model.

Forward Passes and Model Verification

With parameters initialized, we can perform forward passes using the apply() method. This method takes the parameter dictionary and input data, executing the forward computation we defined in __call__. Semantically, apply() represents Flax's functional approach to neural networks: it treats our module as a pure function and applies it to specific parameter values and input data, maintaining JAX's stateless programming paradigm. Python# Perform forward passes predictions_setup = xor_mlp_setup.apply({'params': params_setup}, xor_X) predictions_compact = xor_mlp_compact.apply({'params': params_compact}, xor_X) Let's examine the predictions from both models to verify they produce identical results: textPredictions from MLP (setup approach): [[0.42703056] [0.43364942] [0.44780588] [0.45057902]] Predictions from MLPCompact (@nn.compact approach): [[0.42703056] [0.43364942] [0.44780588] [0.45057902]] Both approaches produce identical results: True The identical parameter structures and predictions confirm that both setup() and @nn.compact approaches are functionally equivalent. The models haven't been trained yet, so the predictions are essentially random outputs from the initialized network; this demonstrates that our MLP architecture is correctly constructed and ready for training.