JAX是一个由Google DeepMind和Google Research的研究人员开发的开源数值计算库。它提供了与NumPy兼容的API，并具备自动微分、即时（JIT）编译以及GPU/TPU加速功能，使其成为机器学习研究和科学计算的基础工具。JAX于2018年首次发布，现已广泛用于训练神经网络和执行高性能数值实验。
JAX的核心设计围绕可组合的函数变换展开。其主要操作包括用于自动微分的grad、用于编译到加速器的jit、用于向量化的vmap以及用于跨设备并行化的pmap。这些变换可以任意嵌套，使研究人员能够用简洁的代码表达复杂算法。JAX使用XLA（加速线性代数）作为其编译后端，该后端针对AMD、Intel、NVIDIA（未在列表中，但隐含提及）以及Google Cloud TPU优化计算。
历史与发展
JAX起源于2017年Google的研究工作，建立在早期autograd和XLA的基础上。首次公开发布于2018年12月。该项目由Matthew Johnson、Roy Frostig和Alex Wiltschko等研究人员领导，并得到了Google Brain团队的广泛贡献。2020年，JAX成为多个知名库的基础，包括Flax（神经网络库）、Haiku（由DeepMind使用）和Trax。到2023年，JAX已成为Google内部机器学习基础设施的核心组件，为大型语言模型和Transformer等模型提供支持。
主要特性
JAX的自动微分支持前向和反向两种模式，能够高效计算任意Python函数的梯度。jit变换通过XLA将函数编译为机器代码，通常比纯Python实现显著加速。vmap自动在批次维度上向量化操作，消除了手动循环展开的需要。pmap将计算分布到多个设备上，便于数据并行和模型并行训练。JAX还包含一个具有函数式API的随机数生成器，确保在不同硬件配置下的可复现性。
该库与Python生态系统无缝集成，支持标准数据结构并与NumPy互操作。它还提供了jax.numpy模块，该模块镜像NumPy的接口，但在加速器设备上运行。JAX的函数式编程风格（数组不可变，函数无副作用）简化了调试，并支持安全的并行执行。
生态系统与采用
JAX催生了一个丰富的专业库生态系统。Flax和Haiku提供高级神经网络API，而Optax提供优化算法。在科学计算领域，JAX-MD（分子动力学）和JAX-COSMO（宇宙学）等库扩展了其应用范围。包括MIT CSAIL、斯坦福AI实验室和伯克利AI研究在内的主要研究机构使用JAX进行强化学习、概率编程和可微仿真等项目。
在工业界，JAX为Google的生产系统提供支持，包括Google Cloud AI服务的部分组件和Waymo的感知模型。OpenAI也将其用于一些研究项目，而Anthropic主要使用PyTorch。该库在TPU上的性能使其成为训练大规模模型的首选，特别是在生成式AI应用中。
与其他框架的比较
JAX与TensorFlow和PyTorch等机器学习框架竞争。与TensorFlow的静态图方法不同，JAX采用函数式、类似NumPy的风格，许多研究人员认为这更直观。与PyTorch相比，JAX在编译和并行化方面提供了更显式的控制，但由于其函数式约束，学习曲线更陡峭。JAX的jit编译通常比PyTorch的即时执行产生更快的推理速度，但PyTorch的动态图更易于调试。在基准测试中，JAX在GPU工作负载上通常与PyTorch相当或更优，并且在TPU上具有明显优势，而PyTorch原生不支持TPU。
应用与未来方向
JAX被用于从深度学习到人工智能研究的多个领域。它为可微物理引擎、贝叶斯推断工具和优化算法提供支持。近期发展包括通过ROCm支持AMD GPU以及改进CPU性能。JAX团队持续增强自动分片和混合精度训练等功能。截至2024年，JAX仍在积极开发中，社区不断壮大，并定期发布新版本。其设计原则（可组合性、性能和可复现性）使其成为下一代AI研究的关键工具。
参见
参考文献
- JAX官方文档和GitHub仓库（访问于2024年）
- Google Research关于JAX的博客文章（2018-2023年）
- 引用JAX的机器学习和科学计算学术论文（2020-2024年）