optax

收录理由

在 JAX 里训练模型，优化器的重要性常常不亚于网络本身，而 Optax 是多数人最终会选择的工具。它把 Adam、带动量的 SGD、AdaGrad 等标准梯度处理组件做成小巧可组合的积木，而不是一套僵硬的训练循环。你可以自由插入学习率调度，或者把多个梯度变换拼成自定义优化器，完全不必改动原有训练代码。这些实现经过充分测试，且写法贴近标准数学公式，让研究者更容易信任和修改其中的逻辑，无论是做实验还是部署到生产环境都够用。对于已经身处 JAX 生态的开发者而言，它是把梯度更新做对的一个务实基础。