简介
JAX 是一个由 Google Research 主导开发的开源 Python 数值计算库，官网文档托管在 readthedocs 上，提供完整的官方介绍与使用指南。它将自动微分（Autograd）、即时编译（JIT）与高性能数组运算结合起来，为科学计算、机器学习、深度学习、大规模并行训练等领域提供底层支持。JAX 的设计紧密遵循 NumPy 风格 API，因此对于熟悉 NumPy 的开发者而言，上手非常自然。它同时支持在 CPU、GPU 和 TPU 等加速器上运行同一份代码，使得开发者可以轻松实现跨平台高效计算。
特点
JAX 的核心优势在于其强大的数值计算能力和灵活的程序转换功能。首先，JAX 提供自动微分功能，可以对各种 Python 函数自动求导，这对于机器学习模型的梯度计算至关重要。其次，通过 JIT 编译，JAX 可以将 Python 代码编译为高效机器码，以充分利用硬件资源提升执行速度，其背后的编译器基于 XLA（Accelerated Linear Algebra）。再者，JAX 支持向量化（vmap）、并行化（pmap）等高级变换，可以大规模并行计算而无需手动写底层代码。此外，JAX 有着丰富的函数转换机制，允许用户组合不同转换来实现复杂的优化策略，这使它成为科研与工程实践中灵活、高效的工具。
如何使用
要开始使用 JAX，通常需要在 Python 环境中安装 jax 和相应的二进制库 jaxlib。安装完成后，你可以像使用 NumPy 一样使用 jax.numpy 来构建数组和实现计算逻辑，但与此同时可通过 jax.grad、jax.jit、jax.vmap、jax.pmap 等接口实现自动微分、即时编译、自动向量化与并行计算。JAX 文档还提供了入门指南、快速入门示例与高级教程，为新手提供逐步学习路径，例如“Quickstart: How to think in JAX”和“JAX 101” 等模块详细阐述 JAX 的基本使用方式与执行模型。此外，官方文档还包括常见问题（FAQ）、错误调试指南与进阶性能优化策略，帮助开发者解决实际使用过程中的各种技术难题。
价格
JAX 本身是一个完全开源且免费的软件库，你可以自由下载、安装和在任何项目中使用，而无需支付许可费用。作为一个 Apache 2.0 许可的开源项目，JAX 可以在商业或学术环境中自由部署，并且可以根据需要进行修改与扩展。需要注意的是，如果你在云平台上使用 GPU/TPU 等加速器资源，这些硬件的使用可能需要付费，但这与 JAX 库本身无直接价格关联。
常见问题
很多新手在使用 JAX 时会有一些共性疑问，例如为什么使用 jax.jit 后函数行为不同于纯 Python 函数，或者如何理解 JAX 的自动微分机制等。官方 FAQ 页面针对这些问题进行了逐条解释，强调像 jax.jit 这样的转换会在内部编译和缓存函数，从而可能导致函数第一次执行与后续调用表现不同。此外，JAX 的执行模型与 NumPy 虽然类似但在细节上有差异，新用户需要理解 JAX 的变换逻辑、跟踪机制与计算图概念，官方文档中也有具体说明。对于更复杂的用例，例如在多个设备上并行计算或自定义操作规则，官方指南也提供了相应的思路和示例，帮助用户逐步掌握 JAX 的高级功能。
