Flax#
为 JAX 打造的神经网络（Neural Networks）
Flax 为使用 JAX 进行神经网络研究的研究人员和开发人员提供了灵活的端到端用户体验。Flax 使您能够充分利用 JAX 的强大功能。
Flax 的核心是 NNX——一个简化的 API，它使得在 JAX 中创建、检查、调试和分析神经网络变得更加容易。Flax NNX 对 Python 引用语义提供了一流的支持，使用户能够使用常规的 Python 对象来表达他们的模型。Flax NNX 是对先前 Flax Linen API 的演进，它凝聚了多年的经验，旨在提供一个更简单、更友好的 API。
注意
由于大多数 Flax 用户仍然依赖 Flax Linen API，因此在短期内不会弃用它。但是，我们鼓励新用户使用 Flax NNX。请查看为什么选择 Flax NNX，了解 Flax NNX 和 Linen 之间的比较，以及我们推出新 API 的原因。
要将您的 Flax Linen 代码库迁移到 Flax NNX，请先在NNX 基础中熟悉该 API，然后按照演进指南开始迁移。
功能特性#
Pythonic 风格
Flax NNX 支持使用常规的 Python 对象，提供了直观且可预测的开发体验。
简洁
Flax NNX 依赖于 Python 的对象模型，这为用户带来了简洁性并提高了开发速度。
表现力强
Flax NNX 允许通过其过滤器（Filter）系统对模型状态进行细粒度控制。
易于上手
Flax NNX 通过函数式 API，使得将对象与常规 JAX 代码集成变得非常容易。
基本用法#
from flax import nnx
import optax
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)) # eager initialization
optimizer = nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param)
@nnx.jit # automatic state management for JAX transforms
def train_step(model, optimizer, x, y):
def loss_fn(model):
y_pred = model(x) # call methods directly
return ((y_pred - y) ** 2).mean()
loss, grads = nnx.value_and_grad(loss_fn)(model)
optimizer.update(model, grads) # in-place updates
return loss