用途：JAX 是面向加速器的数组计算与程序变换库，用 NumPy 风格代码获得梯度、GPU/TPU 加速、向量化与并行化。适合需要自动微分、高性能数值计算、自定义梯度算法、自动向量化/并行化，以及偏好函数式数值代码的场景；简单 NumPy 操作、依赖原地修改或有状态副作用、需要运行时数据决定控制流、或依赖 NumPy 数组的库则不适合。
主要能力：四个可组合变换——jax.grad（反向模式自动微分，支持高阶导数、pytree、value_and_grad、has_aux、jacfwd/jacrev、custom_vjp/custom_jvp）、jax.jit（编译到 XLA，首次调用含编译开销，后续按形状/类型缓存）、jax.vmap（自动批处理向量化，支持 in_axes 与嵌套）、jax.pmap（跨多设备并行）。变换可任意组合，如 jax.jit(jax.vmap(jax.grad(f)))。
与 NumPy 的关键差异：数组不可变，需用 x.at[idx].set/add/mul/min/max 做函数式更新；随机数用显式 PRNGKey 并 split，不可复用；默认 32 位，可用 config.update("jax_enable_x64", True) 或环境变量 JAX_ENABLE_X64=1 开启 64 位；越界索引会被静默钳制（属未定义行为，不应依赖）；非数组输入需显式转换。
使用方式：安装命令资料未提供。基本用法为 import jax、import jax.numpy as jnp，用 @jax.jit 装饰函数，用 jax.grad 求梯度，用 jax.vmap 批处理。JIT 内需纯函数，控制流用 jax.lax.cond、fori_loop、while_loop，动态形状用 jnp.where，调试用 jax.debug.print，循环累积用 jax.lax.scan。静态参数用 static_argnums 标记，固定形状可避免重复编译。
适用场景：优化循环、小批量训练、RNN/scan、pytree 参数梯度、多 GPU/TPU 数据并行、大规模仿真。
限制：JIT 要求纯函数、静态形状与可追踪控制流；离散操作（如 argmax）梯度为零；.at 返回新数组不修改原数组；print 在 JIT 中仅追踪时输出一次。具体版本、安装方式与性能数据资料未提供。