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

核心能力与快速开始

★ 36307 · Python

能做什么

JAX 可自动微分原生 Python 与 NumPy 函数，支持反向模式（jax.grad）和正向模式微分，且可任意阶组合；可通过 jax.jit 使用 XLA 将函数端到端编译，在 TPU、GPU 等硬件加速器上运行；可通过 jax.vmap 沿数组轴自动向量化函数，将矩阵-向量乘法转为矩阵-矩阵乘法等以提升性能；支持将编译与自动微分任意组合，例如生成高效的 Jacobian 或逐样本梯度；提供编译器自动并行、显式分片与自动分区、手动单设备编程等扩展方式，以便在数千设备上扩展计算。

适用场景

适用于高性能数值计算与大规模机器学习场景，例如需要自动微分、JIT 编译和向量化的科学计算与深度学习程序；需要在 TPU、GPU 等加速器上编译并扩展 NumPy 程序的场景；需要逐样本梯度、Jacobian 矩阵或任意阶导数的研究实现；以及需要跨多设备进行编译自动并行、显式分片或每设备手动编程的大规模训练与计算任务。

项目形态

开发框架

部署方式

README 提供安装方式：CPU 使用 pip install -U jax；NVIDIA GPU 使用 pip install -U "jax[cuda13]"；Google TPU 使用 pip install -U "jax[tpu]"；AMD GPU（Linux）使用 pip install -U "jax[rocm7-local]"；Intel GPU 需参考 Intel 的说明。文档还提及从源码编译、使用 Docker、其他 CUDA 版本和社区支持的 conda 构建等替代安装策略。支持平台包括 Linux x86_64、Linux aarch64、Mac aarch64、Windows x86_64 和 Windows WSL2 x86_64，不同平台对 CPU、NVIDIA GPU、Google TPU、AMD GPU、Apple GPU 和 Intel GPU 的支持程度不同。

主要依赖

需要 Python 环境，安装方式按目标平台区分，可分别安装 CPU、NVIDIA GPU、Google TPU、AMD GPU 或 Intel GPU 版本。README 未说明具体硬件最低配置、内存要求或其他运行依赖。

使用限制

README 明确说明这是研究项目，不是 Google 官方产品，使用时应预期存在 sharp edges（常见陷阱），并链接到“Gotchas Notebook”说明。使用 jax.jit 会限制函数可用的 Python 控制流类型。README 未说明 API 成本、性能上限或其他使用限制。

许可证

Apache-2.0

最近推送

2026-09-17 06:54

资料检查：2026-09-17 07:00 · 卡片整理：2026-09-17 07:00
卡片由 AI 根据仓库资料整理，未进行安装或运行验证 · 本次更新未完成，保留上次成功内容

职业应用

程序员

程序员可用 JAX 编写高性能数值计算与机器学习程序：用 jax.grad 对原生 Python/NumPy 函数自动微分（含反向与正向模式、任意阶导数），用 jax.jit 将纯函数编译到 GPU/TPU 等加速器，用 jax.vmap 自动向量化，并通过编译器自动并行、显式分片或手动单设备编程扩展到多设备。

官方使用入口 ↗

由 AI 根据来源材料整理，未进行安装或运行验证