为 JAX 贡献

为 JAX 贡献#

每个人都可以为 JAX 做出贡献，我们重视每个人的贡献。有几种贡献方式，包括：

JAX 项目遵循 Google 的开源社区指南。

贡献方式#

我们欢迎拉取请求，特别是针对那些标记为 contributions welcome 或 good first issue 的问题。

对于其他提案，我们要求您首先在 GitHub 问题 或 讨论 中寻求关于您计划贡献的反馈。

使用拉取请求贡献代码#

我们所有的开发工作都使用 git 进行，因此假设您具备基础知识。

按照以下步骤贡献代码：

  1. 签署 Google贡献者许可协议 (CLA)。更多信息，请参见下面的拉取请求检查清单。

  2. 通过点击 仓库页面 上的 Fork 按钮来 Fork JAX 仓库。这会在你自己的账户中创建一个 JAX 仓库的副本。

  3. 为了运行测试，请在本地安装 Python >= 3.10。

  4. 从源代码安装你的 pip 分支。这允许你修改代码并立即测试它：

    git clone https://github.com/YOUR_USERNAME/jax
cd jax
pip install -r build/test-requirements.txt  # Installs all testing requirements.
pip install -e ".[cpu]"  # Installs JAX from the current directory in editable mode.

  5. 将 JAX 仓库添加为上游远程仓库，以便你可以使用它来同步你的更改。

    git remote add upstream https://www.github.com/google/jax

  6. 创建一个分支，你将从这里开始开发：

    git checkout -b name-of-change

    并使用您喜欢的编辑器实现您的更改（我们推荐 Visual Studio Code）。

  7. 确保你的代码通过 JAX 的 lint 和类型检查，通过从仓库顶部运行以下命令：

    pip install pre-commit
pre-commit run --all

    更多详情请参见 代码检查和类型检查。

  8. 确保测试通过，请从仓库的顶部运行以下命令：

    pytest -n auto tests/

    JAX 的测试套件非常大，所以如果你知道覆盖你更改的特定测试文件，你可以将测试限制在该文件上；例如：

    pytest -n auto tests/lax_scipy_test.py

    你可以通过使用 pytest -k 标志来匹配特定的测试名称，从而进一步缩小测试范围：

    pytest -n auto tests/lax_scipy_test.py -k testLogSumExp

    JAX 还提供了对运行哪些特定测试的更细粒度控制；更多信息请参见 运行测试。

  9. 一旦你对你的修改满意，按照以下方式创建一个提交（ 如何写提交信息）：

    git add file1.py file2.py ...
git commit -m "Your commit message"

    然后与主仓库同步你的代码：

    git fetch upstream
git rebase upstream/main

    最后，将你的提交推送到你的开发分支，并在你的fork中创建一个远程分支，你可以用它来创建一个拉取请求：

    git push --set-upstream origin name-of-change

    请确保您的贡献是一个单独的提交（参见 单一变更的提交和拉取请求）

  10. 从 JAX 仓库创建一个拉取请求并发送以供审核。在准备您的 PR 时，请查看 JAX 拉取请求检查清单 以了解注意事项，如果您需要更多关于使用拉取请求的信息，请参考 GitHub Help。

JAX 拉取请求检查清单#

在准备 JAX 拉取请求时，请记住以下几点：

Google 贡献者许可协议#

对该项目的贡献必须附有Google贡献者许可协议（CLA）。您（或您的雇主）保留您贡献的版权；这只是授予我们使用和重新分发您贡献作为项目一部分的权限。请访问 https://cla.developers.google.com/ 查看您当前的协议文件或签署新的协议。

通常你只需要提交一次CLA，所以如果你已经提交过一次（即使是为了不同的项目），你可能不需要再提交了。如果你不确定是否签署了CLA，你可以打开你的PR，我们友好的CI机器人会为你检查。

单一变更的提交和拉取请求#

一个 git 提交应当是一个自包含的、单一的变更，并附有描述性的信息。这有助于审查，并且在以后发现问题时，有助于识别或恢复变更。

拉取请求通常包含一个 git 提交。（在某些情况下，例如大型重构或内部重写，它们可能包含多个提交。）在准备拉取请求以供审查时，您可能需要将多个提交压缩在一起。如果可能的话，我们要求您在发送 PR 进行审查之前执行此操作。git rebase -i 命令可能对此有所帮助。

代码检查和类型检查#

JAX 使用 mypy 和 ruff 来静态测试代码质量；在本地运行这些检查的最简单方法是使用 pre-commit 框架：

pip install pre-commit
pre-commit run --all-files

如果你的拉取请求涉及文档笔记本，这也会对这些笔记本运行一些检查（更多详情请参见 更新笔记本）。

完整的 GitHub 测试套件#

您的 PR 将自动在 GitHub CI 上通过一个完整的测试套件运行，该套件涵盖了多种 Python 版本、依赖版本和配置选项。这些测试发现您在本地未捕获的失败是正常的；要修复这些问题，您可以向您的分支推送新的提交。

受限测试套件#

一旦你的PR被审查，JAX维护者将把它标记为 pull ready。这将触发一组更大的测试，包括在GPU和TPU后端上的测试，这些测试在标准的GitHub CI中不可用。这些测试的详细结果不公开可见，但分配给你的PR的JAX维护者将与你沟通任何可能发现的失败；例如，数值测试在TPU上的容差通常需要与CPU不同。

目录