自动微分

JAX：面向加速器的高性能数值计算与可组合函数变换库

jax-ml/jax · Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more

JAX 是一个 Python 库，用于面向加速器的数组计算与程序变换，面向高性能数值计算和大规模机器学习。它可以自动微分原生 Python 与 NumPy 函数，支持循环、分支、递归和闭包中的微分以及任意阶导数，并提供反向模式与正向模式微分。JAX 借助 XLA 将 NumPy 程序编译并扩展到 TPU、GPU 及其他硬件加速器，支持用 jax.jit 编译纯函数，编译与自动微分可任意组合。此外，JAX 提供 jax.vmap 自动向量化，并支持编译器自动并行、显式分片与手动单设备编程等扩展方式。该 README 说明这是一个研究项目，而非 Google 官方产品，使用时应预期一些需要注意的边界问题。