企业官网建设流程全解析

最近一年我把不少模型从PyTorch搬到了JAX，起因也是看了几个项目和论文的官方实现——训练代码简洁到让人怀疑人生，自动微分和编译加速是天然集成的。真正动手写了一段时间之后，我的感受是：JAX的确不是一个"换个框架写同样代码"的替代品，它的函数式范式会倒逼你把计算流程想得比之前更清楚。这篇文章就用一个非常经典的模型——图卷积网络（GCN）——来串一遍JAX的核心用法。内容会从JAX的基本功开始，一直写到两层GCN的完整实现和训练，包含能直接复制运行的代码，以及我在踩坑之后才想明白的几个关键点。

这篇内容适合谁？如果你是刚接触JAX、之前主要用PyTorch或者纯NumPy，想搞清楚jit、grad、vmap到底是什么，以及它们如何组合在一起工作；或者你想在一个真实的小型图数据集上跑通GCN但不想被TensorFlow那套生态绑住，这篇都很合适。读完之后你至少能自己动手复现一个GCN，并且对JAX的编程习惯有一个整体把握。

1. JAX的底层气质：先想清楚它为什么长这样

1.1 它是NumPy，但比NumPy更"纯净"

JAX最表面的形态是jax.numpy，接口长得和NumPy几乎一样。第一次跑jnp.array([1,2,3])的时候你会觉得毫无门槛，随便写两行矩阵运算，手感完全一致。但核心区别藏在"数组不可变"这个规定里。在NumPy或者PyTorch里，你习惯了x[0] = 1原地修改，或者params.grad自动累积梯度，这些在JAX里都不存在——数组一旦创建就不能改，想更新某个位置只能通过.at操作符返回一个新数组。

import jax.numpy as jnp x = jnp.array([1.0, 2.0, 3.0]) # x[0] = 10.0 # 报错，TypeError y = x.at[0].set(10.0) print(x) # [1. 2. 3.]，x没变 print(y) # [10. 2. 3.]，y是新数组

这个设计初看很啰嗦，但它是整个JAX能实现jit和grad的前提。因为程序里没有"共享可变状态"，编译器才能放心大胆地做表达式重写、算子融合、甚至是并行切分。换句话说，JAX选择了一条"信任程序纯度"的路：计算纯粹是由输入到输出的函数映射，所有副作用都被剔除。你写的每一段训练逻辑，最后都会被编译成一张无状态的计算图。

我刚从PyTorch转过来的时候最难受的就是这件事。PyTorch里模型参数就是一个Tensor，你随时可以param.grad累加、optimizer.step()原地更新；到了JAX里，参数就是普通的数据结构，用jax.tree_util管理，更新是"老参数 + 负学习率×梯度"这样一个生成新结构的过程。你会被迫写一个显式的update函数。这件事一开始觉得烦，但后来发现有个实际好处：训练循环是确定性的，不会因为某个状态被意外修改产生隐性问题，复现实验变得异常简单。

1.2grad不是装饰器，而是函数变换器

PyTorch里做反向传播要用loss.backward()，然后参数上的.grad被填好，再交给优化器。JAX里没有这一套隐式状态，它把"求梯度"变成一个显式的函数变换：grad(f)接受一个函数，返回这个函数的梯度版本。

import jax import jax.numpy as jnp def f(x): return jnp.sum(x ** 2) g = jax.grad(f) print(g(jnp.array([1.0, 2.0, 3.0]))) # [2. 4. 6.]

注意几点：grad默认只对第一个参数求导；函数必须返回标量，否则你不知道梯度应该"顺着"哪个方向传。如果你有一个函数要同时对多个参数求导，用argnums指定位置即可。如果某个参数是嵌套结构——比如一个字典的字典——grad也能正常返回相同结构的梯度。这一点非常重要，因为GCN的两个权重矩阵和两个偏置，我通常就装在一个简单的列表里，grad照样能算出结构完全一致的梯度。

另外还有jax.value_and_grad，它在求梯度的同时把函数值也返回回来。训练里我们通常既想看loss又需要梯度，用这个变换就可以只写一遍前向计算，省得forward被调用两次。

loss_and_grad = jax.value_and_grad(loss_fn) loss_val, grads = loss_and_grad(params, x, A_hat, labels, mask)

这里loss_fn的签名是loss_fn(params, ...)，我习惯让params永远放在第一位，这样value_and_grad的默认行为正好指向它，不用额外设置argnums。

1.3jit到底加速了什么

jit是JAX的即时编译器，底层通过XLA把Python函数翻译成高效的、可以融合算子的计算图。用起来就是在函数定义上加一个装饰器，或者对已经定义的函数调用一次jax.jit(func)。

@jax.jit def forward(params, x, A_hat): h = A_hat @ x @ params[0] + params[1] h = jax.nn.relu(h) h = A_hat @ h @ params[2] + params[3] return h

第一次调用这个函数时，XLA会把Python层的计算流程编译成底层指令，这个编译过程通常需要几百毫秒甚至更久；之后相同输入形状的调用都走编译好的代码，速度会有明显提升。我的经验是，如果模型不算太复杂，jit带来的收益在小数据集上不一定直观，但一旦batch size或者图的规模上来了，算子融合和内存分配的优化就非常明显，跑大图的时候能省一半以上的时间。

jit也有自己的脾气，最大的限制是它要求函数的控制流和数据形状最好在编译期就能确定。如果函数内部写if x.sum() > 0这种依赖数组值的Python分支，jit会出错或者表现奇怪。遇到这种情况需要用jax.lax.cond这类带条件语义的原语替代。GCN的前向计算本质是矩阵乘法链，几乎碰不到这种问题，这也是我挑它来做示范的原因之一。

1.4 随机数大概是新人最不习惯的地方

JAX的随机数设计是另一个劝退点。PyTorch里torch.randn直接按全局随机状态生成；JAX则要求你显式携带一个"随机密钥"（key），每次生成随机数之前必须用jax.random.split来"分裂"出新key。否则你会发现两次调用生成的是完全一样的数组——这不是bug，是设计。

key = jax.random.PRNGKey(42) # 需要两个正态数组时： key1, key2 = jax.random.split(key, 2) w1 = jax.random.normal(key1, (input_dim, hidden_dim)) w2 = jax.random.normal(key2, (hidden_dim, output_dim))

这种设计一开始很烦，尤其是写训练循环的时候，模型里一旦有dropout就必须在每轮迭代里重新split一个key，然后把它传进forward。但好处是随机性可以被完全复现：只要保存好初始key，整个实验序列都能重演。GCN的初始化代码我会在第三节完整给出，到时候你会看到所有随机数组都是这样"从key里长出来"的，逻辑非常透明。

2. GCN的前向公式与邻接矩阵归一化原理

2.1 图卷积到底"卷积"了什么

GCN是图神经网络里最经典、最朴素的一支。它的设计动机非常简单：普通卷积网络处理的是规则网格（图像、语音），而图上的节点没有固定顺序，也没有平移不变性，所以我们没法定义"卷积核滑过邻域"这件事。Kipf那篇经典工作给出的替代方案是：每一层先通过邻接矩阵聚合邻居的特征，再乘上一个可学习的权重矩阵。它的核心公式是：

H^(l+1) = σ( Â · H^(l) · W^(l))

其中Â是经过归一化的邻接矩阵，H是节点特征，W是权重。这个公式看起来就是一个矩阵连乘，没有任何复杂的结构，所以用JAX来实现非常自然——jnp.dot和@运算符就能覆盖全部计算。

动态上理解这个公式：邻接矩阵A里A_ij = 1表示节点i和节点j之间有边，那么A·H就是"每个节点把自己邻居的特征拿过来相加"。这个操作的意义是：每个节点的表示不再只看自己，而是看自己和周围邻居的组合。一个节点如果邻居很多，它的聚合结果就很大，数值不稳定；如果邻居很少，又显得更新量很小。为了公平起见，需要用度矩阵D来归一化。

2.2 为什么是 D^(-1/2)·A·D^(-1/2) 而不是 D^(-1)·A

归一化有两种常见做法。第一种是行归一化，D^(-1)A，效果是每行之和为1，所有邻居的特征做加权平均。第二种是Kipf用的对称归一化，D^(-1/2)AD^(-1/2)，效果是每个节点更新时，邻居的贡献会被两端的度同时缩放。后者有非常明确的谱域意义：它对应一个归一化图拉普拉斯算子上的线性滤波器，数学上能推导出它近似一种低通滤波，让相似的节点表示更相似。

从实现角度讲，对称归一化的计算方式也很优雅：

def normalize_adjacency(A): A_tilde = A + jnp.eye(A.shape[0]) deg = jnp.sum(A_tilde, axis=1) deg_inv_sqrt = jnp.power(deg, -0.5) deg_inv_sqrt = jnp.where(jnp.isinf(deg_inv_sqrt), 0.0, deg_inv_sqrt) D_inv_sqrt = jnp.diag(deg_inv_sqrt) return D_inv_sqrt @ A_tilde @ D_inv_sqrt

这里加单位阵的作用是给每个节点加自环。如果不加自环，节点在聚合时看不到自己的特征，只能看到邻居，这样关于自身的信息在每一层都会被稀释，深层模型很容易丢失节点身份信息。加了自环相当于默认"我自己也是我邻居之一"，这被后来的无数实验证实是GCN的标配。jnp.where(jnp.isinf(...))是为了防止孤立节点（度为0）在取-0.5次方时出现无穷大。

你可能注意到jnp.power(deg, -0.5)中的deg是零的时候结果会是无穷大，在转换成对角矩阵后会生成NaN。所以先把它替换成0，既保证了数值安全，也确保孤立节点不会贡献任何聚合信息。这一步在PyTorch里你得费点心思，但在JAX里，jnp.where配合jnp.isinf就是一次优雅的向量化操作。

2.3 GCN层数的"2层就够"直觉

实现GCN之前你还应该知道一个经验结论：大多数标准场景下2层GCN已经够了。原因主要是过平滑问题——图卷积本质上是一种平滑操作，每一层都在让邻居节点的表示趋于一致，层数越多，所有节点的表示越发趋同，反而损害分类精度。所以你在各种论文里看到的基线方案，基本就是2层：输入 -> 隐藏层（ReLU） -> 输出层。用我这个项目做例子：输入维度是特征维度，隐藏层设成16，输出维度是类别数，参数总量大概只有2×（输入维×16 + 16×类别数），小到CPU上跑几百轮都毫无压力。

这个"模型小但效果好"的特质，让它特别适合作为JAX的练手项目——你不需要调用复杂的自定义算子，也不需要处理分布式训练，前向和反向都干净利落，但又能把grad、jit、随机数、参数更新全部串起来。

3. 用JAX实现一个两层GCN的完整过程

3.1 准备数据：图结构、特征和标签

我这里用一个经典的论文引用网络数据集来演示，它的规模很小，适合在CPU上快速验证。数据集包含约2700个节点，每个节点是一篇论文，特征是对应词袋向量（0/1表示是否包含某个关键词），标签是论文的类别；节点之间的边表示引用关系。目标是仅用"一部分节点的标签"来训练模型，然后预测剩下的节点。

由于JAX本身不带数据加载器，预处理部分可以直接用常规Python处理。通常这种数据集会提供三个数组：features（节点特征矩阵）、adjacency（邻接矩阵）、labels（标签向量），还有train_mask/val_mask/test_mask分别标记哪些节点用于训练、验证、测试。

import numpy as np import jax.numpy as jnp # 假设已经从数据文件加载到了numpy数组 # features: (N, F) # adjacency: (N, N) 稀疏或稠密 # labels: (N,) 整数 # train_mask: (N,) bool N = features.shape[0] F = features.shape[1] C = int(labels.max()) + 1 A = np.asarray(adjacency.todense(), dtype=np.float32) # 小图可以直接转稠密 X = jnp.array(features, dtype=jnp.float32) Y = jnp.array(labels, dtype=jnp.int32) A_hat = jnp.array(normalize_adjacency(jnp.array(A, dtype=jnp.float32)), dtype=jnp.float32) train_mask = jnp.array(train_mask, dtype=jbool) val_mask = jnp.array(val_mask, dtype=jbool) test_mask = jnp.array(test_mask, dtype=jbool)

如果你用的是稀疏邻接矩阵，JAX对稀疏的支持目前在CPU上还不够顺手，所以小图我建议直接稠密化。2700×2700的矩阵才几十MB，完全没问题。对大图希望用稀疏实现的话，我后面会简单提一下方向，但本篇文章以理解原理为主。

3.2 初始化参数：JAX风格的参数管理

参数初始化在JAX里有一个需要强迫自己记住的原则：参数必须是纯数据结构，不能挂任何隐藏状态。我把两个层的权重和偏置放在一个简单的列表里：

def init_params(key, input_dim, hidden_dim, output_dim): key1, key2 = jax.random.split(key, 2) w1 = jax.random.normal(key1, (input_dim, hidden_dim)) * jnp.sqrt(2.0 / input_dim) b1 = jnp.zeros((hidden_dim,)) w2 = jax.random.normal(key2, (hidden_dim, output_dim)) * jnp.sqrt(2.0 / hidden_dim) b2 = jnp.zeros((output_dim,)) return [w1, b1, w2, b2]

缩放因子sqrt(2/input_dim)是参考了ReLU网络的初始化经验，目的是让前向传播的方差保持稳定。如果你不缩放，网络在深层时数值很容易爆炸，GCN虽然只有两层，但最好还是从一开始就养成好习惯。

init_params返回的是一个Python列表。JAX的jit、grad、value_and_grad可以自动递归处理列表、字典、元组这些结构，这也是jax.tree_util在背后做的事。你可以放心地把这个列表传给梯度函数，返回的梯度依然是一个同构的列表。

3.3 前向传播：两个矩阵乘法就够

前向传播函数定义为：

def forward(params, X, A_hat): w1, b1, w2, b2 = params h = A_hat @ X @ w1 + b1 h = jax.nn.relu(h) logits = A_hat @ h @ w2 + b2 return logits

这里的三次矩阵乘法分别对应：输入特征在图上的第一次平滑传播（聚合邻居信息）、隐藏层到输出层的第二次平滑传播。你可以观察到这个实现没有显式的"邻居采样""注意力机制""边特征融合"，它就是最朴素的聚合加变换。恰恰因为朴素，它才能在JAX里被jit完美加速——整个计算就是一个纯线性代数表达式，XLA可以无压力地做算子融合。

有一点值得注意：这里A_hat @ X @ w1先做A_hat @ X，再做结果与w1的乘法。从计算顺序上讲，先和图矩阵乘会让聚合发生在特征变换之前，这与一些实现先X @ w1再做聚合在数学上其实等价，因为矩阵乘法满足结合律。但两种顺序的浮点误差和效率可能有细微差异，我在实践中习惯让A_hat在最左边，符合"先聚合再变换"的直觉。

3.4 损失函数与训练循环

GCN训练一般用交叉熵。因为只有带标签的节点才有监督信号，所以要把mask应用进去。

def loss_fn(params, X, A_hat, Y, mask): logits = forward(params, X, A_hat) log_probs = jax.nn.log_softmax(logits) label_onehot = jax.nn.one_hot(Y, num_classes=logits.shape[-1]) loss = -jnp.mean(jnp.sum(log_probs * label_onehot, axis=-1) * mask) return loss

这里的mask是布尔向量，True表示该节点参与loss计算。乘法的效果是让未标记节点贡献为0，jnp.mean分母上会包含0贡献的元素，所以严格来说这不是标准的平均CE；不过在图半监督场景下大家一般都用这个写法，因为它简单，且避免了一些除零判断。如果你想做严格的平均，可以改成loss.sum() / mask.sum()，差别很小。

接下来是训练循环的主角——update函数。它做了四件事：计算loss、计算梯度、更新参数、返回新的loss值。整个过程完全无副作用，每次调用返回新的参数结构。

@jax.jit def update(params, X, A_hat, Y, mask, lr): loss_val, grads = jax.value_and_grad(loss_fn)(params, X, A_hat, Y, mask) new_params = [ p - lr * g for p, g in zip(params, grads) ] return new_params, loss_val

你可能会问：参数更新里用了Python的列表推导式，jit能处理吗？能。因为params和grads都是长度固定的列表，推导式在编译阶段就被展开了，XLA清楚地知道每一步操作是什么。这里lr是一个Python浮点数，它会作为编译期的常量参数被烘焙进图里。所以如果你每轮都改变学习率，会导致JIT重新编译。实际训练中我们的学习率通常是固定值或按epoch调整，不是每步都变，这个问题不大。

完整训练循环写出来大概是：

key = jax.random.PRNGKey(0) params = init_params(key, F, 16, C) lr = 0.01 for epoch in range(200): params, loss_val = update(params, X, A_hat, Y, train_mask, lr) if epoch % 20 == 0: acc = evaluate(params, X, A_hat, Y, test_mask) print(f"epoch {epoch}: loss={loss_val:.4f} test_acc={acc:.4f}")

evaluate函数可以顺便演示一下vmap或者argmax：

def evaluate(params, X, A_hat, Y, mask): logits = forward(params, X, A_hat) preds = jnp.argmax(logits, axis=-1) correct = jnp.sum((preds == Y) * mask) total = jnp.sum(mask) return correct / total

这一套全跑完之后，在标准引文网络数据集的实验上你大概率能看到测试准确率在0.75到0.80之间。不同随机种子和初始化下会有几个点的波动，但整体趋势是稳定的。如果你发现结果差得离谱，第一件事检查有没有把自环加上，第二件事检查学习率是否太大——我见过有人设成0.5直接发散到NaN。

4. 实际训练过程中踩过的坑与排查链路

4.1 JIT的隐式重新编译让训练"越跑越慢"

我第一次用JAX跑图模型时遇到过一个很隐蔽的性能问题：训练前20轮正常，后面突然每一步都变慢，而且CPU风扇狂转。排查了半天发现，问题出在A_hat的形状上——训练集合大小不变，按理说形状不会变。但实际上我的数据预处理把A_hat做成了稠密矩阵，而当时代码里用了np.asarray(adjacency.todense())，todense()在某些稀疏后端上会返回一个奇怪的矩阵形状元信息，导致传给jit函数时每次都被当作新的输入签名，于是每个step都在重新编译。

排查思路建议分三步：

• 第一步，把update函数里的@jax.jit临时去掉，看速度是否正常。如果去掉后恢复正常，基本确定是编译开销。
• 第二步，在update函数外打印A_hat.shape和A_hat.dtype，确认这两者是稳定的。JIT的缓存键基于输入数组的形状和dtype，任何一个变化都会触发重编译。
• 第三步，如果确实有某个输入在反复变化形状，把它固定成统一形状，或者把它移出jit函数，只对内部矩阵乘法做预计算。

这个坑最麻烦的地方在于：JAX不会给你任何报错，它只是安静地重编译，而你觉得程序"没坏"，只是慢得诡异。希望这里写出来能帮你省掉一个晚上的排查时间。

4.2random.key管理失当导致Dropout失效

我在另一个图模型里用了dropout，当时天真地以为像PyTorch一样每个step自动有不同的随机行为。结果训练准确率怎么都上不去，后来才发现：我在每次迭代里把同一个key传进了forward，而jax.random.bernoulli(key, ...)对同一个key永远返回同一个mask，所以dropout实际上变成了一个"固定比例的确定性丢弃"，而非随机丢弃。

正确的做法是用jax.random.split产生新的key，并把新key的一部分用于dropout，另一部分保留给下一次迭代：

@jax.jit def update_with_dropout(params, X, A_hat, Y, mask, key, lr): key, subkey = jax.random.split(key) ... logits = forward(params, X, A_hat, subkey) ...

也就是说，key是训练循环里必须显式流转的状态，就像参数一样。GCN本身没有dropout时不会遇到这个问题，但你一旦开始尝试正则化，这个坑会准时出现。

4.3 过平滑现象：深GCN反而更差

我在验证模型时顺手把层数从2层加到4层、6层，结果测试准确率不升反降。这不是JAX的问题，而是图卷积的天然特性：每多一层，每个节点就更倾向于"吸收"更远邻居的信息，节点表示彼此趋同，最终变成所有节点一个模子。后来我专门看了节点表示的相似度，确实是一层比一层高，这印证了过平滑的判断。

所以在使用GCN时，一个实用的建议是：不要盲目堆深度。如果你的任务需要更多层的感受野，尝试用残差连接、JK-Net、或者GAT这类更复杂的架构，不要指望单纯加深GCN能带来收益。用JAX实现残差连接也很简单，就是把上一层的表示和当前层的输出加在一起再激活，一行代码的事。

4.4 特征矩阵的数值范围

某些数据集的特征是词袋0/1向量，数值范围很小，训练很稳定。但如果你换到带TF-IDF权重的特征，范围差异会很大，前向传播里可能出现极大值，经过softmax后损失变成NaN。遇到这种情况，我的经验是检查A_hat是否真的归一化正确，以及输入特征做一次行归一化或者log变换。JAX里做这个非常自然：

X = X / jnp.linalg.norm(X, axis=1, keepdims=True)

4.5 为什么我在代码里避免使用Python原生循环做邻居聚合

如果你之前在PyTorch里写过GNN，可能会习惯用for循环对每个节点手动找邻居。在JAX里，这种方式在jit下很难高效编译，因为循环次数取决于节点度，是运行时才确定的值，XLA通常无法把动态循环展开成高效指令。正确的姿势是把整个图的操作向量化成矩阵乘法——用A_hat @ H一次性计算所有节点的聚合结果。这正是GCN公式本身已经做好的事。

如果将来你需要实现GraphSAGE那种邻居采样算法，推荐的方向是用jax.lax.scan替代Python的for循环，或者在jit外部用Python循环逐批处理。用vmap也可以对批量节点应用同一个聚合函数，这是JAX进行图模型优化的通用工具箱。

5. 进一步优化：JIT加速可视化、vmap、以及更真实的训练体验

5.1 用block_until_ready看到真实的JIT收益

JAX是异步执行的，你写loss_val = update(...)的时候，计算可能还没真正跑完，只是被放进了队列。如果你用Python的time模块测速，测到的可能是"提交时间"而不是"实际计算时间"，这会让性能测试产生严重误判。解决办法是在计时前调用jax.block_until_ready。

import time start = time.time() for _ in range(100): params, loss_val = update(params, X, A_hat, Y, train_mask, lr) loss_val.block_until_ready() elapsed = time.time() - start print(f"100 steps: {elapsed:.2f}s")

第一次调用update时会包含JIT编译时间，所以你在测速时最好先跑一次"预热"，再正式开始计时。我在比较jit开和关的性能差距时踩过这个坑，第一次跑出来的时间大得离谱，后来才知道是编译开销。

5.2vmap在GCN场景中的实际价值

虽然GCN的前向本身是全局矩阵乘法，不需要vmap参与，但在模型调优和诊断阶段，vmap非常有用。比如你想对同一组参数快速测试10个不同随机种子的初始化效果，可以把初始化和前向都映射到一个向量化函数上：

from functools import partial def init_and_forward(key, X, A_hat): params = init_params(key, F, 16, C) return forward(params, X, A_hat) keys = jax.random.split(jax.random.PRNGKey(0), 10) logits_batch = jax.vmap(partial(init_and_forward, X=X, A_hat=A_hat))(keys)

vmap自动把keys的第一个维度当成batch维，其他参数保持不变。它能让你在写批量实验时避免手写Python循环，而且这种向量化方式在jit下通常能获得更好的编译结果。这个技巧在你做超参数扫参的时候非常实用，建议收藏。

5.3 训练过程的动态日志与早停

我还想分享一个很实在的经验：哪怕模型很小，也建议写一个基于验证集准确率的早停机制。GCN的loss曲线在标准引文数据集上经常是先快速下降，然后在一个平台期轻微波动，如果没有早停，最后取到的checkpoint未必是验证集上最好的。JAX没有现成的EarlyStopping回调，但实现起来只需要维护一个"最优参数"的副本和计数变量：

best_params = params best_val_acc = 0.0 patience = 20 wait = 0 for epoch in range(500): params, loss_val = update(params, X, A_hat, Y, train_mask, lr) val_acc = evaluate(params, X, A_hat, Y, val_mask) if val_acc > best_val_acc: best_val_acc = val_acc best_params = params wait = 0 else: wait += 1 if wait >= patience: break

这个模式在几乎所有JAX项目里都能复用。它的价值在于：避免你拿测试集反复调参造成信息泄露，同时在模型接近过平滑时及时停下来，而不是等到所有节点都变成同一个向量。

5.4 模型评估的一个细节：mask的数据类型

最后提醒一个容易忽略的细节：mask数组的dtype必须是布尔型。如果你从数据集里读到的mask是整数0/1，jnp.sum(mask)也能算，但一旦参与乘法，整型mask会把浮点梯度变成全零，或者意外截断浮点数。我在一次实验里就因为这个把测试准确率算成了0，排查了半天。稳妥的做法是读取数据后统一做一次类型转换：

mask = jnp.array(mask, dtype=jbool)

6. 一些额外的个人体会

从PyTorch迁到JAX，最值回票价的不是某个具体算子的快慢，而是一种思维上的"瘦身"。你被迫把训练过程拆成init、forward、loss、update这么几个纯函数，任何中间状态都无处躲藏，实验逻辑因此变得极度透明。GCN作为一个只需十行前向代码的模型，恰恰能让你把注意力完全放在JAX本身的特性上，而不是被复杂的网络结构干扰。

如果你已经顺利跑通了上面的代码，我建议你接着做三件事：一是把hidden_dim从16调整到64、128，观察训练速度和精度的变化；二是把dropout加进两层之间，用jax.random.split正确管理key，看看正则化对最终准确率的影响；三是把数据集换成一个更大规模的引文网络，感受一下稠密A_hat带来的内存压力，这时你可以开始研究JAX对稀疏矩阵的支持，或者用vmap配合邻居采样来实现更灵活的训练方式。这三步做完，你对JAX和GCN的掌握都会再上一个台阶。