4、JAX模型基础入门：Flax的nn.Module、Haiku的hk.Module、JAX的纯函数式编程范式

说实话，刚接触JAX生态的时候，我也有点懵。

PyTorch里一个nn.Module走天下，到了JAX这边，突然冒出Flax、Haiku、还有纯函数式编程。到底该选哪个？它们之间有什么区别？

这一章，我就带你把这些东西理清楚。

4.1 为什么JAX需要「框架」？

JAX本身其实是个数值计算库，类似NumPy的升级版。它不提供「模型层」、「优化器」、「训练循环」这些东西。

你想想看，如果让你用纯JAX写一个神经网络，你得手动管理所有参数、手动写前向传播、手动处理随机种子。写个简单的线性层还好，一旦到了ResNet、Transformer这种级别，代码量会爆炸。

所以，社区搞出了Flax和Haiku。它们的作用就是：在JAX的纯函数式基础上，给你提供一套好用的模型构建工具。

4.2 Flax的nn.Module：最像PyTorch的写法

我个人习惯用Flax。为什么？因为它上手最快。

你看下面这个例子，定义一个简单的MLP：

import flax.linen as nn
import jax.numpy as jnp

class SimpleMLP(nn.Module):
    hidden_dim: int
    output_dim: int

    @nn.compact
    def __call__(self, x):
        x = nn.Dense(self.hidden_dim)(x)
        x = nn.relu(x)
        x = nn.Dense(self.output_dim)(x)
        return x

# 初始化
model = SimpleMLP(hidden_dim=128, output_dim=10)
rng = jax.random.PRNGKey(0)
x = jnp.ones((32, 64))
params = model.init(rng, x)
y = model.apply(params, x)

看到没？nn.Module、@nn.compact、model.init、model.apply——这些概念和PyTorch几乎一一对应。

Flax的init返回的是一个参数字典，而不是模型对象。这是JAX的函数式核心：模型和参数是分离的。你调用model.apply(params, x)时，把参数传进去，模型才真正运行。

4.3 Haiku的hk.Module：更显式的参数管理

Haiku的写法稍微不同。它要求你在__call__里显式声明参数，而不是像Flax那样自动注册。

import haiku as hk
import jax.numpy as jnp

def forward_fn(x):
    mlp = hk.Sequential([
        hk.Linear(128),
        jax.nn.relu,
        hk.Linear(10),
    ])
    return mlp(x)

# 转换函数
forward = hk.transform(forward_fn)

# 初始化
rng = jax.random.PRNGKey(0)
x = jnp.ones((32, 64))
params = forward.init(rng, x)
y = forward.apply(params, None, x)

注意这里的关键：hk.transform。它把一个普通函数「转换」成带参数管理的函数。我个人觉得，Haiku这种方式更「JAX原生」——因为它强迫你理解「函数 + 参数」的分离。

4.4 纯JAX函数式编程：从零手写一个线性层

为了让你真正理解JAX的哲学，我们来看纯JAX怎么写一个线性层。

import jax
import jax.numpy as jnp

def init_linear(rng, in_dim, out_dim):
    """初始化线性层的参数"""
    key_w, key_b = jax.random.split(rng)
    w = jax.random.normal(key_w, (in_dim, out_dim)) * 0.01
    b = jax.random.normal(key_b, (out_dim,)) * 0.01
    return {'w': w, 'b': b}

def forward_linear(params, x):
    """前向传播"""
    return x @ params['w'] + params['b']

# 使用
rng = jax.random.PRNGKey(42)
params = init_linear(rng, 64, 10)
x = jnp.ones((32, 64))
y = forward_linear(params, x)

看到了吗？没有类，没有对象，只有函数和字典。这就是JAX的纯函数式编程范式：

• 参数是显式的：你手动创建、手动传递
• 没有隐藏状态：所有状态都在参数里
• 函数是纯的：同样的输入，永远得到同样的输出

你想想看，这种设计有什么好处？

好处太多了。比如：

• 容易调试：参数就是字典，你可以随时打印、保存、修改
• 容易并行：纯函数天然适合vmap、pmap等并行变换
• 容易优化：JAX的grad、jit等变换可以直接应用

4.5 三者的对比与选择

我整理了一个表格，方便你对比：

特性	Flax	Haiku	纯JAX
上手难度	低（像PyTorch）	中	高
参数管理	自动注册	通过transform	手动管理
灵活性	中	中高	最高
社区生态	Google官方，生态好	DeepMind，学术圈多用	无框架依赖
适合场景	快速原型、生产部署	研究、需要精细控制	底层优化、教学

4.6 本章知识体系

下面这张图，帮你梳理本章的核心逻辑：

嗯，这张图把三者的关系说得很清楚了。无论你用Flax、Haiku还是纯JAX，底层都是JAX的函数式核心——参数和模型分离，函数无状态。

4.7 本章小结

这一章我们聊了：

• Flax的nn.Module：自动参数注册，写法最像PyTorch
• Haiku的hk.Module：通过hk.transform转换，参数管理更显式
• 纯JAX函数式：手动管理参数，最灵活但也最繁琐
• 三者的对比和选择建议

说实话，没有「最好」的框架，只有「最适合」你的框架。我个人建议你先从Flax入手，等熟悉了JAX的思维模式后，再尝试其他方式。

下一章，我们会真正开始动手——用Flax和PyTorch分别实现同一个模型，然后进行互转。到时候你会发现，理解了本章的内容，后面的互转操作会变得非常自然。

---

公众号:蓝海资料掘金营，微信deep3321