jax

用途：JAX 面向高性能数值与科学计算、深度学习研究，通过 Autograd 自动微分与 XLA 加速线性代数，支持 GPU/TPU。主要能力：jit 即时编译、grad 自动求导、vmap 自动批处理向量化、pmap 与 shard_map 多设备分布式训练；可搭配 Flax、Haiku 或 Equinox 构建神经网络。使用方式：jax.numpy 提供类 NumPy 接口，核心函数为纯函数，用 @jax.jit 编译性能关键函数，jax.grad 求梯度，jax.vmap 对单样本函数按批次自动向量化（如 in_axes=(None,0)）。随机数需显式管理：用 jax.random.PRNGKey 生成根键，每次随机操作前用 jax.random.split 拆分，不可复用同一键。限制与注意：函数应无副作用、数组不可原地修改，需用 x.at[idx].set(val) 形式更新；JIT 内不可用依赖动态数组值的 Python if/while，应改用 jax.lax.cond 或 jax.lax.while_loop；不要在变换函数内修改全局变量。常见错误包括 ConcretizationTypeError（改用 lax.cond）、数组不可变赋值失败（改用 .at[].set()）、PRNG 键复用导致随机数相同（先 split）。