JAX

JAX 是由 Google Brain 团队开发的高性能数值计算库，提米哥愿意称之为“打了类固醇的 NumPy”。对于普通用户来说，它的 API 看起来和 NumPy 几乎一模一样，但其背后隐藏着巨大的能量。JAX 的核心是两大“黑科技”：自动微分（automatic differentiation）和 XLA（Accelerated Linear Algebra）编译。自动微分意味着 JAX 可以对任何原生的 Python 和 NumPy 函数进行求导，这对于训练神经网络至关重要。而 XLA 编译器则可以将你的 Python/NumPy 代码即时编译（JIT）成在 GPU 和 TPU 上高效运行的优化代码。这两者的结合，使得 JAX 兼具了 NumPy 的易用性和底层框架的极致性能。它深受全球顶尖 AI 研究机构（尤其是 DeepMind）的喜爱，很多突破性的研究成果，比如 AlphaFold 2（蛋白质结构预测），都是基于 JAX 开发的。相比于 PyTorch，JAX 在处理函数式编程、模型并行化和复杂变换方面给予了研究者更大的自由度和灵活性。虽然它的学习曲线比 PyTorch 陡峭，生态也相对较小，但对于那些追求极致性能、需要进行前沿算法探索的硬核研究员来说，JAX 是他们手中最锋利的一把剑。