贡献 JAX

贡献 JAX#

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

在 JAX 的讨论区回答问题
改进或扩展 JAX 的文档
为 JAX 的代码库做出贡献
为构建在 JAX 之上的更广泛的库生态系统做任何上述贡献

JAX 项目遵循 Google 的开源社区行为准则。

贡献方式#

我们欢迎拉取请求，特别是对于标记为“欢迎贡献”或“第一个好问题”的 issue。

对于其他提案，我们要求您首先打开一个 GitHub Issue 或 Discussion，以征求对您计划贡献的反馈。

通过拉取请求贡献代码#

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

请按照以下步骤贡献代码：

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

从源代码 pip 安装您的 fork。这使您可以修改代码并立即进行测试。

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/jax-ml/jax
```
创建一个分支来开发您的更改。
```
git checkout -b name-of-change
```
使用您喜欢的编辑器实现您的更改（我们推荐 Visual Studio Code）。
通过在仓库根目录运行以下命令，确保您的代码通过 JAX 的代码风格和类型检查：
```
pip install pre-commit
pre-commit run --all
```
有关更多详细信息，请参阅代码风格和类型检查。
通过在仓库根目录运行以下命令，确保测试通过：
```
pytest -n auto tests/
```
同时在 64 位模式下运行它们，方法是设置环境变量 JAX_ENABLE_X64=True。
```
JAX_ENABLE_X64=True 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
```
最后，将您的提交推送到您的开发分支，并在您的 fork 中创建一个远程分支，您可以使用该分支来创建拉取请求。
```
git push --set-upstream origin name-of-change
```
请确保您的贡献是单次提交（请参阅单次更改的提交和拉取请求）。
从 JAX 仓库创建拉取请求并提交审查。请查看 JAX 拉取请求清单以了解准备 PR 时的注意事项，如果您需要有关使用拉取请求的更多信息，请查阅 GitHub 帮助。

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。这将触发一组更大的测试，包括在标准 GitHub CI 无法使用的 GPU 和 TPU 后端上的测试。这些测试的详细结果不可公开查看，但分配给您 PR 的 JAX 维护者会与您沟通可能发现的任何失败；例如，数值测试在 TPU 上的容差可能与在 CPU 上不同，这种情况并不少见。