nnx_module

Chapter 1: nnx.Module

Welcome to the Flax NNX tutorial! This first chapter introduces the cornerstone of building neural networks with NNX: nnx.Module.

Motivation: Pythonic Neural Networks in JAX

JAX provides powerful tools for accelerated numerical computing, particularly automatic differentiation and compilation (XLA). However, its functional nature can sometimes make defining and managing the state (parameters, batch statistics, etc.) of complex neural networks cumbersome.

Traditional Flax (Linen) addressed this with a functional API built around nn.Module, which required special methods (setup, compact) and external state management.

Flax NNX introduces nnx.Module to offer a more intuitive, object-oriented approach. The core idea is simple: define your network layers and models as standard Python classes. State is held directly as instance attributes, initialization happens in __init__, and forward computation is defined in regular methods (like __call__). This leverages familiar Python principles, simplifying debugging, inspection, and overall development workflow, while still integrating seamlessly with JAX's functional transformations through helper APIs.

Central Use Case: Defining a Simple Linear Layer

Let's see how nnx.Module works by defining a basic linear transformation layer.

import jax
import jax.numpy as jnp
from flax import nnx
from flax.nnx.nn import initializers # For parameter initialization

# Define needed RNGs for initialization
rngs = nnx.Rngs(0)

class SimpleLinear(nnx.Module):
  def __init__(self, in_features: int, out_features: int, *, rngs: nnx.Rngs):
    # Define state (parameters) directly as attributes
    # We use nnx.Param, a type of nnx.Variable, to mark learnable parameters
    self.kernel = nnx.Param(
        initializers.lecun_normal()(rngs.params(), (in_features, out_features))
    )
    self.bias = nnx.Param(initializers.zeros_init()(rngs.params(), (out_features,)))

  def __call__(self, x: jax.Array) -> jax.Array:
    # Access state attributes directly in forward computation
    y = jnp.dot(x, self.kernel.value) + self.bias.value
    return y

# Instantiate the layer like a regular Python class
layer = SimpleLinear(in_features=3, out_features=4, rngs=rngs)

# Create dummy input data
x = jnp.ones((1, 3))

# Call the layer instance directly
y = layer(x)

print(f"Output shape: {y.shape}")
# Output: Output shape: (1, 4)