JAX 是一个将 NumPy 风格 API 与可组合函数变换结合的框架，核心变换包括 Grad（微分）、Jit（编译）、Vmap（向量化）和 Pmap（并行化），可把 Python/NumPy 代码通过 XLA 编译优化到 CPU、GPU 或 TPU。
用途：适合需要 GPU/TPU 加速的高性能科学仿真、自定义机器学习研究、计算高阶导数（Hessian、Jacobian）、物理信息机器学习与可微分仿真、函数自动向量化，以及同一套代码跨设备运行。
主要能力：jax.numpy 镜像 NumPy 但支持变换；jax.grad 求导、jax.hessian 求二阶导；jax.jit 即时编译；jax.vmap 自动批处理；jax.pmap 多设备并行；jax.random 手动管理 PRNG 状态；变换可作用于 PyTree（字典、列表、元组等嵌套数组容器）。
使用方式：安装命令为 pip install jax jaxlib（CPU），GPU 为 pip install "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html，具体 CUDA 版本需查文档。常用导入为 import jax、import jax.numpy as jnp、from jax import grad, jit, vmap, pmap, random。
限制与规则：函数必须是纯函数，不能有副作用；数组不可变，需用 x = x.at[idx].set(val) 做函数式更新；JIT 中避免 Python 控制流，改用 jax.lax.cond、while_loop、fori_loop；随机数必须手动拆分 key，如 key, subkey = random.split(key)；JIT 中非数组参数需用 static_argnums 标记；异步执行下计时需 block_until_ready()；GPU 默认预分配 90% 显存，可用 XLA_PYTHON_CLIENT_PREALLOCATE 环境变量调整。建议先在 CPU 调试再扩展到 GPU/TPU，并注意首次 JIT 编译较慢。