JAX 是什么？
JAX 是一个面向加速器的 Python 数值计算库，专为高性能数组计算与大规模机器学习设计，支持在 CPU、GPU 和 TPU 等多种硬件后端上高效执行相同代码。它提供类 NumPy 的熟悉 API，大幅降低研究人员与工程师的学习门槛；内置可组合的函数变换能力，包括即时编译（jit）、自动微分（grad）、向量化（vmap）和并行化（pmap），使开发者能以简洁表达式构建高性能计算图；所有变换均基于纯函数语义与不可变数组，确保可预测性与跨平台一致性。JAX 本身轻量专注，其生态则由 Flax、Optax、Equinox 等专业库协同扩展，覆盖神经网络、优化器、微分方程求解与数据加载等全栈需求。
主要功能
提供类 NumPy 的 API，便于研究人员和工程师快速上手
内置可组合的函数变换：包括编译（jit）、自动微分（grad）、向量化（vmap）和并行化（pmap）
适用场景与人群
面向研究人员与工程师的Python库，适用于需要高性能数组计算、可编程变换及大规模机器学习的开发者
- 用于高性能数值计算，如大规模机器学习模型训练与科学仿真
- 实现自动微分、向量化、即时编译与并行化等可组合函数变换
- 跨CPU、GPU与TPU后端运行同一套代码，支持多设备与分布式计算
Python函数、NumPy风格数组、随机密钥、pytree结构化数据等
经变换后的JAX数组（DeviceArray）、编译后可执行对象（Compiled）、梯度函数、向量化函数等
常见问题
JAX是否提供自己的神经网络框架？
JAX本身不提供高层神经网络框架，但与Flax、Equinox、Keras等生态库集成，推荐通过JAX AI Stack构建模型
JAX能否在不同硬件后端（如CPU/GPU/TPU）上运行同一代码？
可以，JAX设计目标之一即‘Run anywhere’，同一代码可在CPU、GPU与TPU上原生执行
如何使用
- 01
安装JAX库：pip install jax[jaxlib]
- 02
导入并使用jax.numpy替代NumPy进行数组运算
- 03
应用函数变换如jax.grad()实现自动求导或jax.jit()加速执行
- 04
在CPU、GPU或TPU上运行相同代码，无需修改逻辑
官网与资料来源
资料更新于 2026-10-30
