为 JAX 贡献代码

内容

为 JAX 贡献代码#

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

回答 JAX 讨论页面上的问题
改进或扩展 JAX 的文档
为 JAX 的代码库做出贡献
以上述任何方式为更广泛的基于 JAX 的库生态系统做出贡献

JAX 项目遵循 Google 开源社区准则.

贡献方式#

我们欢迎拉取请求，特别是对于那些标有欢迎贡献或良好入门问题的问题。

对于其他提议，我们要求您首先打开一个 GitHub 问题或讨论以寻求对您计划的贡献的反馈。

使用拉取请求贡献代码#

我们使用 git 进行所有开发，因此假设您具备基本知识。

请按照以下步骤贡献代码

签署Google 贡献者许可协议 (CLA)。有关更多信息，请参阅下面的拉取请求清单。
通过点击存储库页面上的 **Fork** 按钮来分叉 JAX 存储库。这会在你自己的帐户中创建一个 JAX 存储库的副本。
在本地安装 Python >= 3.10 以便运行测试。

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.

将 JAX 存储库添加为上游远程存储库，以便你使用它同步更改。
```
git remote add upstream https://www.github.com/google/jax
```
创建一个你将从中开发的分支。
```
git checkout -b name-of-change
```
并使用你喜欢的编辑器实现你的更改（我们推荐Visual Studio Code）。
通过从存储库的顶部运行以下命令，确保你的代码通过了 JAX 的 linting 和类型检查。
```
pip install pre-commit
pre-commit run --all
```
有关更多详细信息，请参阅Linting 和类型检查。
通过从存储库的顶部运行以下命令，确保测试通过。
```
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 还提供了对运行哪些特定测试的更细粒度的控制；有关更多信息，请参阅运行测试。
一旦你对你的更改满意，请按照以下步骤创建一个提交（如何编写提交消息）。
```
git add file1.py file2.py ...
git commit -m "Your commit message"
```
然后将你的代码与主存储库同步。
```
git fetch upstream
git rebase upstream/main
```
最后，将你的提交推送到你的开发分支，并在你的分叉中创建一个远程分支，你可以使用该分支创建拉取请求。
```
git push --set-upstream origin name-of-change
```
请确保你的贡献是一个单一提交（请参阅单一更改提交和拉取请求）。
从 JAX 存储库创建一个拉取请求并将其发送以供审查。在准备你的 PR 时，请查看JAX 拉取请求清单，并咨询GitHub 帮助，如果你需要有关使用拉取请求的更多信息。

JAX 拉取请求清单#

在准备 JAX 拉取请求时，以下是一些需要牢记的事项。

Google 贡献者许可协议#

对该项目的贡献必须附带 Google 贡献者许可协议 (CLA)。你（或你的雇主）保留对你贡献的版权；这只是允许我们在项目中使用和重新分发你的贡献。前往https://cla.developers.google.com/查看你目前存档的协议或签署新的协议。

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

单一更改提交和拉取请求#

一个 git 提交应该是一个独立的、单一更改，带有描述性消息。这有助于审查，并帮助在以后发现问题时识别或撤消更改。

**拉取请求通常包含一个 git 提交。**（在某些情况下，例如对于大型重构或内部重写，它们可能包含多个。）在准备拉取请求以供审查时，你可能需要将多个提交压缩在一起。我们要求你尽可能在发送 PR 以供审查之前这样做。命令git rebase -i 在这方面可能会有用。

Linting 和类型检查#

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 上不同的容差，这并不罕见。