贡献 JAX#
每个人都可以为 JAX 做出贡献,我们重视每个人的贡献。 有几种贡献方式,包括
回答 JAX 的讨论页面上的问题
改进或扩展 JAX 的文档
为 JAX 的代码库做出贡献
以上述任何方式为更广泛的基于 JAX 构建的库生态系统做出贡献
JAX 项目遵循Google 的开源社区准则。
贡献的方式#
我们欢迎拉取请求,特别是对于那些标有欢迎贡献或适合新手的问题。
对于其他提案,我们要求您首先打开一个 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/
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(即使它是针对不同的项目),您可能不需要再次提交。 如果您不确定是否已签署 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 上不同的容差。