jax

JAX 为 Python 和 NumPy 程序提供了一套可组合的函数转换系统。它旨在通过使用 XLA (Accelerated Linear Algebra) 进行编译，将数值计算扩展到包括 GPU 和 TPU 在内的硬件加速器上。 核心能力包括： - 自动微分：通过 `jax.grad`，它支持在 Python 循环、分支、递归和闭包中进行反向模式和前向模式微分，从而实现高阶导数。 - 即时 (JIT) 编译：通过 `jax.jit`，它利用 XLA 编译纯函数，以提高在加速器上的执行性能。 - 自动向量化：通过 `jax.vmap`，它将函数沿数组轴映射，通过将循环下推至原语操作，消除了手动管理批处理维度的需要。 - 扩展与并行化：支持多种扩展模式，包括基于编译器的自动并行化、显式分片以及使用显式集合的手动单设备编程。 JAX 兼容多个平台，包括 Linux、macOS 和 Windows，支持 CPU、NVIDIA GPU、Google TPU、AMD GPU，并对 Apple 和 Intel GPU 提供实验性支持。