JAX 是一个 Python 库，用于面向加速器的数组计算和程序变换，面向高性能数值计算与大规模机器学习。它结合了熟悉的 NumPy 风格 API 与自动微分、编译、向量化和并行化等函数变换能力。核心价值：编写类似 NumPy 的 Python 代码，无需改变数学记号，即可通过可组合的函数变换自动获得梯度、GPU/TPU 加速、向量化与并行化。
适用场景：需要自动微分做优化或机器学习；希望以极少代码改动获得 GPU/TPU 加速；需要高性能数值计算；构建自定义的基于梯度的算法；需要自动向量化或并行化函数；研究场景需要灵活微分；偏好函数式风格的数值代码。
不适用场景：无性能需求的简单 NumPy 操作（开销不划算）；大量依赖原地修改（JAX 数组不可变）；命令式、有状态、带副作用的代码；控制流依赖运行时数据值（支持有限）；依赖必须使用 NumPy 数组的库（存在兼容问题）。
核心变换（可组合）：
- jax.grad：反向模式自动微分，返回计算梯度的函数，可组合求高阶导数，支持嵌套结构（pytrees）。
- jax.jit：将函数编译到 XLA 以获得显著加速。首次调用包含编译开销较慢，相同形状/dtype 的后续调用使用缓存。
- jax.vmap：沿批维度自动向量化函数，消除手写循环，代码更声明式。
- jax.pmap：跨多设备（GPU/TPU）并行，用于多设备计算、训练中的数据并行和大规模模拟。 变换可无缝组合，例如 jit(grad(f))、vmap(grad(f))、jit(vmap(grad(f)))；组合顺序影响性能但不影响正确性。
与 NumPy 的关键差异：
- 不可变性：JAX 数组不可原地修改，使用 x.at[idx].set/add/mul/min/max 返回新数组，支持多索引与切片。
- 随机数：使用显式 PRNGKey，需先 split 再使用，不可复用密钥；可用 vmap 生成并行随机数。
- 64 位精度：默认 32 位，可通过 config.update("jax_enable_x64", True) 或环境变量 JAX_ENABLE_X64=1 启用。
- 越界索引：JAX 静默钳制而非报错（未定义行为，不应依赖）。
- 非数组输入：JAX 要求显式 jnp.array 转换，避免追踪期间性能下降。
易踩坑点：要求纯函数（副作用不可靠）；值依赖的控制流需用 jax.lax.cond，循环用 jax.lax.fori_loop / while_loop；输出形状不能依赖运行时值（用 jnp.where 代替掩码索引）；.at 返回新数组不改原数组；JIT 内 print 只在追踪时打印一次，应使用 jax.debug.print；argmax 等离散操作无梯度，可用 softmax 等可微近似。
自动微分能力：jax.grad 求梯度，jax.value_and_grad 同时返回值与梯度，argnums 指定参数，has_aux 返回辅助数据；支持高阶导数、jax.hessian、jax.jacfwd/jacrev 求雅可比；可用 jax.custom_vjp / custom_jvp 自定义梯度。
JIT 使用：装饰器或 jax.jit(fn)；循环次数等需用 static_argnums 标记静态参数；不同形状会触发重编译，应保持形状一致。
PyTree：字典等嵌套容器可整体求梯度，支持 tree.map、tree.reduce、tree_flatten/unflatten。
常见模式：优化循环（jit(grad(loss)) 后按学习率更新参数）、小批量训练步、用 jax.lax.scan 实现高效累积求和与 RNN 等循环计算。
限制：原文未提供该技能条目的安装命令、版本要求、收费或平台信息，暂未提供。