ABOUT THIS TOOL
关于 JAX
面向高性能数值计算的开源 Python 库，提供 NumPy 风格数组、自动微分、编译、批处理和并行变换。它面向“构建可微分并可在 CPU、GPU 或 TPU 上执行的数值程序”这一类科研任务，可通过Python、CPU、GPU、TPU、Windows、macOS、Linux使用。本站将其归入数据分析，相关能力标签包括自动微分、加速计算。
适合把数值计算写成可组合的数组函数，再按需要使用自动微分、即时编译、批处理和并行变换。开发时应先用小规模输入验证函数、形状与梯度，再记录数据类型、后端、设备、随机键、编译设置以及 JAX 和 jaxlib 版本。
JAX 适合哪些科研任务？
- 构建可自动微分的数值模型
- 用 JIT 编译和向量化加速数组计算
- 在 CPU、GPU 或 TPU 上运行并行计算
选择与使用建议
建议先用不含敏感信息的小型任务试用，检查输出是否可追溯、能否导出，以及是否适合你所在团队的工作流。涉及论文结论、统计分析或专业模型时，应保留原始来源、参数和人工复核记录。
- JAX 变换对纯函数、数组形状和跟踪期控制流有明确约束，普通 Python 副作用可能不按直觉执行
- 硬件后端、精度设置和异步执行会影响性能与数值行为；Windows、macOS 和不同加速器的支持范围并不相同
自动微分加速计算
最后核验：2026-09-09
资料来源：官方来源 ↗
工具功能、价格与政策可能变化，请以官方页面为准。