贡献 JAX#

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

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

贡献方式#

我们欢迎 pull request,特别是对于那些标记为 contributions welcomegood first issue 的问题。

对于其他提案,我们要求您首先打开 GitHub IssueDiscussion,以寻求对您计划贡献的反馈。

使用 pull request 贡献代码#

我们所有的开发都使用 git,因此假定您具备基本知识。

按照以下步骤贡献代码

  1. 签署 Google 贡献者许可协议 (CLA)。 有关更多信息,请参阅下面的 Pull Request 检查清单。

  2. 通过单击 存储库页面上的 Fork 按钮来 Fork JAX 存储库。 这将在您自己的帐户中创建 JAX 存储库的副本。

  3. 在本地安装 Python >= 3.10 以运行测试。

  4. 使用 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.
    
  5. 将 JAX 仓库添加为上游远程仓库,以便您可以使用它来同步您的更改。

    git remote add upstream https://www.github.com/jax-ml/jax
    
  6. 创建一个您将在其上进行开发的 branch

    git checkout -b name-of-change
    

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

  7. 通过从存储库的顶部运行以下命令,确保您的代码通过 JAX 的代码检查和类型检查

    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. 一旦您对您的更改感到满意,请按如下所示创建一个 commit( 如何编写 commit 消息

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

    然后将您的代码与主仓库同步

    git fetch upstream
    git rebase upstream/main
    

    最后,将您的 commit 推送到您的开发 branch 上,并在您的 fork 中创建一个远程 branch,您可以从中创建 pull request

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

    请确保您的贡献是单个 commit(请参阅 单次更改提交和 pull request

  10. 从 JAX 存储库创建一个 pull request 并发送以供审核。 在准备您的 PR 时,请查看 JAX pull request 检查清单 以了解注意事项,如果您需要有关使用 pull request 的更多信息,请查阅 GitHub 帮助

JAX pull request 检查清单#

在您准备 JAX pull request 时,以下是一些需要牢记的事项

Google 贡献者许可协议#

对此项目的贡献必须附带 Google 贡献者许可协议 (CLA)。 您(或您的雇主)保留对您贡献的版权;这仅仅是授予我们允许将您的贡献用作项目一部分并重新分发的权限。 请访问 https://cla.developers.google.com/ 查看您当前的文件协议或签署新的协议。

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

单次更改提交和 pull request#

一个 git commit 应该是一个自包含的、带有描述性消息的单次更改。 这有助于审核,并有助于在稍后发现问题时识别或还原更改。

Pull request 通常包含单个 git commit。 (在某些情况下,例如对于大型重构或内部重写,它们可能包含多个。)在准备用于审核的 pull request 时,您可能需要将多个 commit 合并在一起。 我们要求您尽可能在发送 PR 进行审核之前执行此操作。 git rebase -i 命令可能对此有用。

代码检查和类型检查#

JAX 使用 mypyruff 来静态测试代码质量;在本地运行这些检查的最简单方法是通过 pre-commit 框架

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

如果您的 pull request 涉及到文档 notebook,这也将对它们运行一些检查(有关更多详细信息,请参阅 更新 notebook)。

完整的 GitHub 测试套件#

您的 PR 将自动通过 GitHub CI 上的完整测试套件运行,该套件涵盖一系列 Python 版本、依赖项版本和配置选项。 这些测试出现您在本地未捕获的失败是正常的;要修复这些问题,您可以将新的 commit 推送到您的 branch。

受限测试套件#

一旦您的 PR 经过审核,JAX 维护人员会将其标记为 pull ready。 这将触发更大范围的测试,包括在标准 GitHub CI 中不可用的 GPU 和 TPU 后端上的测试。 这些测试的详细结果不公开可见,但分配给您的 PR 的 JAX 维护人员将与您沟通这些测试可能发现的任何失败;例如,数值测试在 TPU 上比在 CPU 上需要不同的容差是很常见的。