Flax NNX
Google让JAX拥有Python灵魂的神经网络库

🔥从Linen到NNX：一场范式革命

2020年，Google Brain（现Google DeepMind）发布了Flax Linen API，为JAX提供了一套基于函数式语义的神经网络库。Linen的核心设计是lazy初始化——Module在构造时不创建参数，参数只在init方法被调用时通过追踪（trace）生成，然后以嵌套字典的形式返回。前向推理则通过apply方法注入参数执行。

这个设计与Haiku的hk.transform异曲同工，都与JAX的纯函数变换模型完美对齐。在规模化训练和性能优化方面，Linen表现出色：变量集合（collections）的组织、自动PRNG管理、SPMD注解等功能让大规模分布式训练变得高效。

但函数式语义的代价是非Pythonic的编程体验：

调试困难：Module构造后，子模块属性不可访问——block.linear会抛出AttributeError，因为linear只在setup()或compact上下文中才被创建。你无法在构造后检查模型的参数形状或结构。

init/apply双轨制：所有顶层计算必须通过init或apply方法执行。代码被割裂为"apply内部"（可以直接调用方法和子模块）和"apply外部"（必须通过apply间接调用）两个世界。

状态管理复杂：添加一个BatchNorm层意味着你需要处理batch_stats集合，在apply时配置mutable参数，手动管理更新后的状态。每加一种有状态层，训练代码就要改一次。

模型手术困难：由于lazy初始化，替换子模块不保证可行；参数结构与Module结构分离，手动同步容易出错。

2024年，Flax团队给出了答案：Flax NNX。核心思想是引入引用语义——Module是真正的Python对象，持有自己的参数，可以被直接检查、修改和调用。这不是Linen的补丁，而是一次从底向上的重新设计。

🏗️nnx.Module：Python对象持有自己的参数

NNX的Module与Linen的Module有本质区别：NNX Module在构造时就创建并持有参数。这意味着你可以在任何时刻检查、修改和调用Module，无需init/apply的间接层。

from flax import nnx
import jax
import jax.numpy as jnp

class Linear(nnx.Module):
    def __init__(self, din: int, dout: int, *, rngs: nnx.Rngs):
        self.w = nnx.Param(rngs.params.uniform((din, dout)))
        self.b = nnx.Param(jnp.zeros((dout,)))
        self.din, self.dout = din, dout

    def __call__(self, x: jax.Array):
        return x @ self.w + self.b[None]

# 构造即初始化——参数立即可用
model = Linear(2, 5, rngs=nnx.Rngs(params=0))
y = model(x=jnp.ones((1, 2)))  # 直接调用，无需apply

print(model.w)  # Param(value=Array(shape=(2, 5), dtype=float32))
print(model.b)  # Param(value=Array(shape=(5,), dtype=float32))

这段代码展示了NNX的三个核心特征：

Eager初始化：构造函数__init__中创建的nnx.Param立即被赋值。参数是Module的真实属性，不是延迟生成的影子。代价是你必须显式提供输入输出维度——没有shape inference，但这换来的是可预测性和可调试性。

直接调用：model(x)就是前向传播。没有apply，没有参数字典注入。Module自己持有参数，调用时直接使用。

nnx.Param变量类型：参数不是裸的ndarray，而是被包装在nnx.Param中。Param是nnx.Variable的子类，实现了所有数值运算符——你可以直接用self.w参与矩阵乘法，无需self.w.value。但当你需要修改值时，使用.value属性或[...]索引。

与Linen对比：在Linen中，同样的代码需要setup()方法、init()调用和apply()调用，且构造后无法访问self.linear。在NNX中，一切都是Python对象的行为——构造、访问、调用，如同使用PyTorch或Keras。

🎲nnx.Rngs：显式但优雅的PRNG管理

JAX使用显式的PRNG密钥，这在大型模型中意味着大量的密钥分裂和管理。Linen通过隐式的RNG序列自动管理，但代价是调试时难以追踪密钥流向。NNX选择了显式但简洁的方案：

# 单一Rngs对象管理所有随机数源
rngs = nnx.Rngs(0)  # 主密钥

# 不同用途的密钥流
params_rng = rngs.params()      # 参数初始化
dropout_rng = rngs.dropout()    # Dropout
noise_rng = rngs.noise()        # 自定义噪声

# 在Module构造时传入
model = MLP(2, 64, 3, rngs=nnx.Rngs(0))

# 在前向传播时传入新的Rngs（如Dropout需要）
y = model(x, rngs=nnx.Rngs(1))

nnx.Rngs是一个容器对象，持有多个命名的RNG流。每次调用rngs.params()或rngs.dropout()时，它从对应的流中派生一个新的密钥。这比Linen的隐式管理更透明——你可以清楚地看到每个随机操作使用了哪个密钥流。

对于需要多个独立密钥的场景（如多层Dropout），NNX提供了fork方法：

# 为5层模型fork出5个独立的Rngs
param_rngs = nnx.Rngs(0).fork(split=5)
dropout_rngs = nnx.Rngs(1).fork(split=5)

# 每层有自己的密钥流
@nnx.vmap(in_axes=0, out_axes=0)
def create_model(rngs):
    return MLP(10, 32, 10, rngs=rngs)

model = create_model(param_rngs)  # 5个独立的MLP

fork将一个Rngs分裂为多个独立的Rngs，每个都有自己独立的密钥序列。这在nnx.scan和nnx.vmap中特别有用——每层或每个并行分支获得自己的密钥流，确保随机行为的可复现性和正确性。

📦有状态计算：Variable与原地更新

在JAX的纯函数世界中，原地更新（in-place update）是禁忌。但NNX通过引用语义和自定义变换，安全地支持了原地状态更新。这是NNX相对于Linen最大的体验提升之一。

class Count(nnx.Variable): pass

class Counter(nnx.Module):
    def __init__(self):
        self.count = Count(jnp.array(0))

    def __call__(self):
        self.count[...] += 1  # 原地更新

counter = Counter()
print(f'{counter.count[...] = }')  # count = 0
counter()
print(f'{counter.count[...] = }')  # count = 1

这里self.count[...] += 1看起来像原地修改，但在NNX的变换系统内部，它被正确地处理为函数式状态更新。当这段代码在nnx.jit内部运行时，NNX会追踪所有Variable的修改，并在变换边界处将修改同步回原始对象。

一个更实用的例子是自定义BatchNorm：

class BatchNorm(nnx.Module):
    def __init__(self, features: int, mu: float = 0.95):
        self.scale = nnx.Param(jnp.ones((features,)))
        self.bias = nnx.Param(jnp.zeros((features,)))
        self.mean = nnx.BatchStat(jnp.zeros((features,)))
        self.var = nnx.BatchStat(jnp.ones((features,)))
        self.mu = mu  # 静态值，不是Variable

    def __call__(self, x):
        mean = jnp.mean(x, axis=-1)
        var = jnp.var(x, axis=-1)
        # EMA更新——原地修改BatchStat
        self.mean.value = self.mu * self.mean + (1 - self.mu) * mean
        self.var.value = self.mu * self.var + (1 - self.mu) * var
        # 归一化和缩放
        x = (x - mean) / jnp.sqrt(var + 1e-5)
        return x * self.scale + self.bias

注意nnx.BatchStat——它是nnx.Variable的另一个子类，与nnx.Param同级但语义不同。NNX的Filter系统可以区分不同类型的Variable：优化器只更新Param，BatchStat由前向传播自动更新，序列化时可以选择性保存。这种类型驱动的状态管理比Linen的字符串集合（"params"、"batch_stats"）更安全、更Pythonic。

⚡nnx.jit与自动状态传播

NNX最令人印象深刻的设计是自动状态传播。在Linen中，状态管理是手动的——你需要告诉apply哪些集合是mutable的，然后从返回值中提取更新后的状态。在NNX中，这一切都是自动的：

import optax

model = MLP(2, 16, 10, rngs=nnx.Rngs(0))
optimizer = nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param)

@nnx.jit  # 自动状态管理
def train_step(model, optimizer, x, y, rngs):
    def loss_fn(model: MLP, rngs: nnx.Rngs):
        y_pred = model(x, rngs)
        return jnp.mean((y_pred - y) ** 2)

    loss, grads = nnx.value_and_grad(loss_fn)(model, rngs)
    optimizer.update(model, grads)  # 原地更新
    return loss

x, y = jnp.ones((5, 2)), jnp.ones((5, 10))
loss = train_step(model, optimizer, x, y, rngs)

这段代码中有两个关键过程在自动发生：

1. BatchNorm和Dropout的状态自动传播。在loss_fn内部，model(x, rngs)会更新BatchNorm的BatchStat并消耗Dropout的RNG。这些修改发生在nnx.value_and_grad内部，然后通过nnx.jit自动传播到外部的model对象。你不需要手动提取和重新注入状态——NNX的变换系统追踪所有Variable修改，在变换边界处同步回原始引用。

2. Optimizer持有model的引用。optimizer.update(model, grads)直接更新model的参数——这之所以能工作，是因为optimizer在构造时就持有model的引用，这个引用关系在nnx.jit变换内部被保持。你不需要像Linen那样从grads手动更新params字典，然后重新创建apply的输入。

nnx.jit是jax.jit的超集——它接受Module和Variable作为输入，在编译时将它们分解为纯函数可以处理的pytree，在执行后将结果重新组装回对象引用。这种"对象 ↔ pytree"的双向变换是NNX的核心技术。

🔧Functional API：nnx.split与nnx.merge

NNX虽然以引用语义为核心，但仍然提供了与JAX纯函数世界交互的桥梁：Functional API。这个API由三个方法组成：nnx.split、nnx.merge和nnx.update。

class Count(nnx.Variable): pass

class StatefulLinear(nnx.Module):
    def __init__(self, din: int, dout: int, *, rngs: nnx.Rngs):
        self.w = nnx.Param(rngs.uniform((din, dout)))
        self.b = nnx.Param(jnp.zeros((dout,)))
        self.count = Count(jnp.array(0, dtype=jnp.uint32))

    def __call__(self, x: jax.Array):
        self.count.value += 1
        return x @ self.w + self.b

model = StatefulLinear(3, 5, rngs=nnx.Rngs(0))

# 分解：Module → GraphDef + State
graphdef, state = nnx.split(model)

# State是纯pytree，可以传给jax.jit、jax.grad等
# GraphDef是模块的静态结构描述

# 重组：GraphDef + State → Module
model2 = nnx.merge(graphdef, state)
assert model2.w.value is not None  # 参数完整恢复

nnx.split将Module分解为两部分：GraphDef（图的静态结构——类、属性名、静态值）和State（所有Variable的值，以嵌套Mapping形式组织）。State是一个纯pytree，可以被jax.jit、jax.grad等标准JAX变换处理。

nnx.merge是split的逆操作——将GraphDef和State重新组装为Module对象。重组后的Module与原始Module在结构上完全一致，参数值相同。

这个API的核心价值在于在需要时切换语义模型。当你需要与纯JAX代码交互（如自定义变换、序列化、分布式同步）时，用split切换到pytree语义；当你需要面向对象操作时，用merge切换回引用语义。NNX的变换系统（nnx.jit、nnx.grad等）内部就是通过split/merge实现的。

Filter系统允许你在split时选择性地提取特定类型的Variable：

# 只提取Param（可训练参数）
graphdef, params = nnx.split(model, nnx.Param)

# 提取Param和BatchStat
graphdef, (params, batch_stats) = nnx.split(
    model, (nnx.Param, nnx.BatchStat)
)

# 只更新Param
nnx.update(model, params)

这种类型驱动的Filter比Linen的字符串键（['params', 'batch_stats']）更安全——类型检查在编译时完成，拼写错误会被IDE和类型检查器捕获。

🧬模型手术：LoRA与nnx.iter_graph

模型手术（model surgery）——替换、修改、插入子模块——在Linen中是出了名的困难。在NNX中，它就是Python对象操作，因为Module是普通Python对象，属性可以直接赋值。

class LoraParam(nnx.Param): pass

class LoraLinear(nnx.Module):
    def __init__(self, linear: Linear, rank: int, rngs: nnx.Rngs):
        self.linear = linear
        self.A = LoraParam(rngs.normal((linear.din, rank)))
        self.B = LoraParam(rngs.normal((rank, linear.dout)))

    def __call__(self, x: jax.Array):
        return self.linear(x) + x @ self.A @ self.B

# 创建模型
rngs = nnx.Rngs(0)
model = MLP(2, 32, 5, rngs=rngs)

# 模型手术——直接替换子模块
model.linear1 = LoraLinear(model.linear1, 4, rngs=rngs)
model.linear2 = LoraLinear(model.linear2, 4, rngs=rngs)

# 现在model的linear1和linear2是LoRA层
y = model(x=jnp.ones((3, 2)), rngs=rngs)

就这么简单——model.linear1 = LoraLinear(...)。原始Linear的参数被保留在LoraLinear的self.linear中，LoRA的A和B矩阵作为LoraParam存储。你可以通过Filter系统区分原始参数和LoRA参数，对它们使用不同的学习率或优化策略。

对于批量模型手术，NNX提供了nnx.iter_graph：

rngs = nnx.Rngs(0)
model = Block(rngs)

# 遍历模型图，替换所有Linear层
for path, module in nnx.iter_graph(model):
    if isinstance(module, nnx.Linear):
        # 找到父模块并替换
        parent, attr = path[-1]
        setattr(parent, attr, LoraLinear(module, rank=5, rngs=rngs))

在Linen中实现同样的操作需要intercept_methodsAPI——一个不直观且容易出错的方法。在NNX中，它就是一个图遍历加属性赋值，任何Python开发者都能理解和修改。

🔄nnx.scan与nnx.vmap：层堆叠与并行

NNX的变换系统不仅复刻了JAX的变换，还针对神经网络场景做了增强。nnx.scan和nnx.vmap是两个典型例子。

# 用nnx.vmap创建5个独立的MLP
@nnx.vmap(in_axes=0, out_axes=0)
def create_model(rngs):
    return MLP(10, 32, 10, rngs=rngs)

# 用nnx.scan迭代应用每一层
@nnx.scan(in_axes=(0, 0, nnx.Carry), out_axes=nnx.Carry)
def forward(model: MLP, rngs: nnx.Rngs, x):
    x = model(x, rngs)
    return x

# 创建5层MLP堆叠
param_rngs = nnx.Rngs(0).fork(split=5)
model = create_model(param_rngs)

# 前向传播——每层独立处理
x = jnp.ones((3, 10))
dropout_rngs = nnx.Rngs(1).fork(split=5)
y = forward(model, dropout_rngs, x)

这段代码展示了NNX变换的几个独特能力：

1. nnx.vmap可以创建Module堆叠。通过对构造函数做vmap，一次性创建5个独立的MLP对象。每个MLP有自己的参数和状态，存储在一个"堆叠"的Module中。

2. nnx.scan自动传播BatchNorm状态。在scan的每一步中，BatchNorm的BatchStat被自动更新并传播到下一步。这在Linen中需要手动管理mutable状态。

3. nnx.scan有意偏离jax.lax.scan。NNX的scan模仿vmap的API风格，支持多个输入、每个输入/输出的scan轴指定，以及carry位置的控制。这比jax.lax.scan的(f, init, xs)三元组更灵活。

4. Rngs.fork为每层分裂独立密钥流。nnx.Rngs(1).fork(split=5)为5层Dropout生成5个独立的RNG流，确保每层的随机行为独立且可复现。

🏋️训练循环：nnx.Optimizer与optax

NNX提供了内置的nnx.Optimizer，与Google的optax优化器库无缝集成：

from flax import nnx
import optax
import jax
import jax.numpy as jnp

class Model(nnx.Module):
    def __init__(self, din, dmid, dout, rngs: nnx.Rngs):
        self.linear = nnx.Linear(din, dmid, rngs=rngs)
        self.bn = nnx.BatchNorm(dmid, rngs=rngs)
        self.dropout = nnx.Dropout(0.2, rngs=rngs)
        self.linear_out = nnx.Linear(dmid, dout, rngs=rngs)

    def __call__(self, x):
        x = nnx.relu(self.dropout(self.bn(self.linear(x))))
        return self.linear_out(x)

model = Model(2, 64, 3, rngs=nnx.Rngs(0))
optimizer = nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param)

@nnx.jit
def train_step(model, optimizer, x, y):
    loss_fn = lambda model: ((model(x) - y) ** 2).mean()
    loss, grads = nnx.value_and_grad(loss_fn)(model)
    optimizer.update(model, grads)  # 原地更新
    return loss

# 训练循环
for step in range(1000):
    x, y = get_batch()
    loss = train_step(model, optimizer, x, y)
    if step % 100 == 0:
        print(f'Step {step}, loss: {loss}')

print(f'Optimizer steps: {optimizer.step.value}')
# 检查参数
print(f'Final kernel shape: {model.linear.kernel.shape}')
print(f'BatchNorm mean: {model.bn.mean.value}')

注意几个关键设计：

wrt=nnx.Param：Optimizer只更新Param类型的Variable。BatchStat、CacheStat等不会被优化器修改。这比Linen中手动指定可训练参数集合更安全。

optimizer.update(model, grads)：一行代码完成梯度更新。Optimizer持有model的引用，直接修改model的参数。不需要从grads字典手动更新params字典，不需要重新组装apply的输入。

训练后直接检查：model.linear.kernel.shape和model.bn.mean.value可以直接访问。在Linen中，这需要从params和batch_stats字典中提取，且BatchNorm的运行时统计量不在params中。

nnx.value_and_grad：这是jax.value_and_grad的NNX版本，支持Module作为输入。它内部通过split将Module分解为pytree，计算梯度后再通过merge重组。对用户完全透明。

⚖️Linen vs NNX：迁移指南

Flax团队明确表示Linen不会立即废弃——大量生产代码依赖Linen，短期内不会强制迁移。但新项目推荐使用NNX。以下是关键迁移对照：

Module定义：

# Linen
class Block(nn.Module):
    train: bool
    def setup(self):
        self.linear = nn.Dense(10)
        self.bn = nn.BatchNorm(use_running_average=not self.train)
    def __call__(self, x):
        return nn.relu(self.bn(self.linear(x)))

# NNX
class Block(nnx.Module):
    def __init__(self, rngs):
        self.linear = nnx.Linear(5, 10, rngs=rngs)
        self.bn = nnx.BatchNorm(10, rngs=rngs)
    def __call__(self, x):
        return nnx.relu(self.bn(self.linear(x)))

关键差异：Linen用setup()定义子模块（lazy），NNX用__init__（eager）。Linen需要train: bool属性控制BatchNorm行为，NNX的BatchNorm自动管理运行时状态。Linen的nn.Dense需要输入维度推断，NNX的nnx.Linear需要显式指定din和dout。

初始化与调用：

# Linen
model = Block(train=True)
params = model.init(random.key(0), x)['params']
y = model.apply({'params': params}, x)

# NNX
model = Block(nnx.Rngs(0))
y = model(x)

Linen需要init + apply两步，NNX构造即初始化、直接调用。

状态管理：

# Linen——手动管理batch_stats
vs = model.init(key, x)
params, batch_stats = vs['params'], vs['batch_stats']
y, updates = model.apply(
    {'params': params, 'batch_stats': batch_stats},
    x, mutable=['batch_stats']
)
batch_stats = updates['batch_stats']

# NNX——自动管理
model = Block(nnx.Rngs(0))
y = model(x)  # batch_stats自动更新

这是NNX最大的体验提升——添加有状态层不需要修改训练代码。

迁移策略：Flax团队建议先熟悉NNX Basics，然后按模块逐步迁移。NNX和Linen可以在同一代码库共存，不需要一次性全部迁移。对于大型项目，可以新模块用NNX、旧模块保持Linen，逐步替换。

💡生态与未来

Flax不仅仅是一个Module库，它是JAX神经网络研究的完整生态系统。

Optax：梯度处理和优化器库，由DeepMind维护。提供Adam、SGD、LAMB等优化器，以及梯度裁剪、权重衰减、学习率调度等组合式梯度变换。NNX的nnx.Optimizer直接使用Optax的优化器作为输入。

Chex：JAX测试和断言库，提供数据结构验证、参数检查、伪随机测试工具。

Treescope：JAX pytree的可视化库，NNX的nnx.display基于Treescope实现，可以交互式地展开和检查Module的嵌套结构。

Orbax：检查点和序列化库，支持异步保存/恢复、分布式同步、多格式存储。Flax的检查点API基于Orbax构建。

JAX生态集成：Flax与JAX的核心变换（jit、grad、vmap、pmap、scan、remat）深度集成。NNX的变换（nnx.jit、nnx.grad、nnx.vmap、nnx.scan）是JAX变换的超集，在保持引用语义的同时兼容所有JAX功能。

Hugging Face集成：transformers库对Flax有一等支持，许多模型（BERT、GPT、T5等）都有Flax实现。随着NNX的推广，越来越多的模型将迁移到NNX API。

与Haiku的关系：2023年DeepMind推荐新项目使用Flax而非Haiku。NNX的发布进一步巩固了这个方向——NNX的引用语义比Haiku的transform模式更接近研究者的直觉，同时保持了JAX的函数式核心。Haiku进入维护模式，NNX成为DeepMind推荐的JAX神经网络库。

TensorFlow/Keras用户的迁移路径：NNX的eager初始化和直接调用模式让Keras用户感到熟悉。nnx.Module的__init__ + __call__模式与tf.keras.layers.Layer的__init__ + call模式高度相似。PyTorch用户也会发现NNX的nnx.Module与torch.nn.Module在概念上对应——两者都是持有参数的Python对象，都支持直接调用。

📄项目信息

@software{flax2020github,
  author = {Jonathan Heek and Anselm Levskaya and Avital Oliver
            and Marvin Ritter and Bertrand Rondepierre
            and Andreas Steiner and Marc van Zee},
  title = {{F}lax: A neural network library and ecosystem for {JAX}},
  url = {http://github.com/google/flax},
  version = {0.12.8},
  year = {2024},
}