Flax调试：对参数进行哈希处理

- **JAX** 是 Google 开发的数值计算库，采用函数式编程风格：数据不可变，每个操作返回新数组而非修改原数据。这使其在 GPU/TPU 上自动并行和加速非常高效，但也要求代码写法与传统框架不同。 - **Flax** 是基于 JAX 的神经网络库，其 **NNX** API 是较新的接口，试图模拟 PyTorch 那种"就地更新参数"的编程体验，但在 JAX 纯函数式底层上实现这种"副作用"需要特殊处理。 - 本文作者正在用 JAX/Flax 从零训练一个小型语言模型。调试时发现 loss 始终不降（稳定在 10.82，对应模型随机猜测的概率），怀疑梯度更新出了问题。 - 核心问题：作者在训练函数上用了标准的 `@jax.jit`（JAX 的即时编译装饰器），但 NNX 的就地参数更新机制要求使用 Flax 专属的 `@nnx.jit` 才能正确传播状态变化。用错装饰器导致梯度计算正确但参数从未被实际更新。 - 调试技巧：面对 7700 万个参数（不可能肉眼对比数值），作者利用 Python 的 `hash()` 对参数数组的字节表示取哈希——参数若有任何微小变化，哈希值会完全不同，从而快速验证参数是否在更新。