1. 自动微分（jax.grad）：支持反向模式（反向传播）与正向模式求导，可对任意阶数求导，且能透过循环、分支、递归与闭包等Python控制流进行微分。
2. 即时编译（jax.jit）：通过XLA将纯函数端到端编译为高效内核，实现算子融合，显著提升计算性能。
3. 自动向量化（jax.vmap）：把“逐个样本”的计算自动批量化为向量运算，免去手工管理batch维度的繁琐。
4. 大规模并行（jax.pmap / SPMD）：支持在多个GPU/TPU上进行单程序多数据（SPMD）并行编程，可扩展至数千设备。
5. NumPy风格API（jax.numpy）：接口与NumPy几乎一致，老用户可无缝迁移、上手门槛低。
6. 可组合变换：grad、jit、vmap、pmap等变换可任意嵌套组合，灵活构建复杂数值程序。
7. 丰富生态配套：Flax（神经网络模块）、Optax（优化器）、Orbax（checkpoint管理）、Equinox等，构成完整的深度学习工具链。

使用JAX非常简单，主要分为三步：
第一步，安装。CPU版本直接执行 pip install -U jax；如需NVIDIA GPU加速，可安装 pip install -U "jax[cuda13]"；在Google Cloud TPU上则使用 pip install -U "jax[tpu]"。
第二步，导入并编写数值函数。在代码中执行 import jax 与 import jax.numpy as jnp，即可像使用NumPy一样进行数组运算。
第三步，施加变换。用 jax.grad 对函数求梯度、用 @jax.jit 装饰函数以编译加速、用 jax.vmap 自动向量化，例如 grad_loss = jax.jit(jax.grad(loss)) 即可一步得到“编译后的梯度函数”。初学者可直接参考官方Quickstart与“Thinking in JAX”系列教程快速上手。

JAX的核心优势可归结为四点：一是函数式、无状态的编程风格，数组不可变、随机数需显式传入key，代码更纯净、易于复现与推理；二是性能卓越，借助XLA编译与算子融合，JAX在GPU尤其是Google TPU上表现出色，被Anthropic、Apple、NVIDIA及Google等公司用于训练前沿大模型；三是“全能变换”一站式集成，求导、编译、向量化、并行化四大能力天然可组合，极大简化了从算法原型到大规模分布式训练的路径；四是生态日趋成熟，以Flax、Optax、Equinox、Orbax为代表的生态库覆盖模型构建、优化、序列化等全流程，使JAX既能做科研原型，也能支撑生产级AI系统。

在主流深度学习框架中，JAX与PyTorch、TensorFlow常被放在一起比较。与PyTorch相比：PyTorch采用面向对象、命令式动态图模型，开发调试体验友好、生态庞大，是通用研究与工业应用的主流选择；而JAX采用函数式、不可变数组的设计，代码更简洁且易于并行扩展，在高性能数值计算与TPU大规模训练场景中表现更优。与TensorFlow相比：TensorFlow更偏企业级全流程与生产部署，工具链完备；JAX则更轻量、更专注科研与计算效率，二者在Google内部甚至互补使用。总体而言，PyTorch易用性更强、社区更大，TensorFlow部署生态成熟，而JAX在极致性能、函数式简洁性与大规模并行方面独具优势，适合对算力与扩展性有高要求的用户。

Q1：JAX的随机数为什么和NumPy不一样？
JAX没有全局随机状态，随机数通过显式传递的key（jax.random.key）生成，以保证纯函数性与可复现性。生成多个独立随机流时，需用jax.random.split不断派生新key。
Q2：为什么jax.jit编译后某些函数报错？
jit要求函数为“纯函数”，其中不宜使用动态Python控制流（如依赖数据的if/while）或就地修改数组；此类逻辑可用jax.lax.cond、jax.lax.scan或jax.numpy的函数替代。
Q3：JAX与NumPy完全兼容吗？
大部分兼容，但JAX数组不可变（不支持就地赋值、视图写入等），且部分操作受编译约束，迁移时需注意差异。
Q4：JAX是Google的官方正式产品吗？
JAX是Google开源的科研项目，而非官方正式产品，因此可能存在“sharp edges”，社区鼓励用户反馈问题并共同改进。