jax-ml/jax-ai-stack

这是什么？

核心功能

• 01一条命令安装 JAX 及 Flax、Optax、Orbax、Chex、Grain、ml_dtypes，版本经过集成测试验证。
• 02提供 jit、vmap、grad 等 JAX 程序变换。
• 03通过 pip 额外依赖可选安装 tensorflow-datasets。
• 04支持安装面向 CUDA GPU 或 TPU 的硬件专用 JAX 版本。

适用场景

• 搭建基于 Flax 与 Optax 训练神经网络的 JAX 开发环境。
• 在 JAX 项目中使用 Orbax 保存和恢复模型检查点。
• 在 NVIDIA GPU 或 Google TPU 上运行 JAX 任务，并使用经过验证的依赖组合。

快速上手

1. 1

   从 PyPI 安装 JAX AI 技术栈

   pip install jax-ai-stack
2. 2

   可选：通过 tfds 扩展安装 TensorFlow Datasets 支持

   pip install jax-ai-stack[tfds]
3. 3

   可选：在同一条命令中加入 JAX CUDA 扩展以支持 GPU

   pip install jax-ai-stack "jax[cuda]"

使用前须知

• !它是一个打包了其他库的集合，除所固定的包之外没有独立功能。
• !README 称文档和教程仍在完善中。

收录于专题

标签

常见问题

jax-ai-stack 是什么？

一键安装的 JAX 工具包，打包了经过集成测试、版本兼容的 Flax、Optax、Orbax 等 AI 与数值计算库。 JAX AI Stack 是一个统一安装入口，用于安装一组相互兼容的 JAX 生态库。执行 pip install jax-ai-stack 即可同时安装 JAX 以及 Flax（构建神经网络）、Optax（梯度处理与优化）、Orbax（检查点保存）、Chex（编写可靠的 JAX 代码）、Grain（数据加载）和 ml_dtypes（NumPy 数据类型扩展）等组件，各版本已通过仓库中的集成测试验证可以协同工作。

jax-ai-stack 可以免费使用吗？用的什么许可证？

jax-ai-stack 以 Apache-2.0 许可证开源。用于商业产品前，请先确认许可证条款。

jax-ai-stack 还在维护吗？

活跃维护中。jax-ai-stack 最近一次提交是在 2026年10月13日。

jax-ai-stack 适合用来做什么？

搭建基于 Flax 与 Optax 训练神经网络的 JAX 开发环境。 在 JAX 项目中使用 Orbax 保存和恢复模型检查点。 在 NVIDIA GPU 或 Google TPU 上运行 JAX 任务，并使用经过验证的依赖组合。

jax-ai-stack 怎么安装和运行？

从 PyPI 安装 JAX AI 技术栈: pip install jax-ai-stack 可选：通过 tfds 扩展安装 TensorFlow Datasets 支持: pip install jax-ai-stack[tfds] 可选：在同一条命令中加入 JAX CUDA 扩展以支持 GPU: pip install jax-ai-stack "jax[cuda]"

在哪里获取 jax-ai-stack？

源代码和安装说明在 GitHub 上：https://github.com/jax-ml/jax-ai-stack