JAX 增强提案 (JEPs)#
大多数更改都可以通过简单的 issue/讨论和 pull request 进行讨论。
但是,一些更改范围较大或需要更多讨论,这些更改应该作为 JEP 来实现。这样可以编写更长的文档,以便在 pull request 中进行讨论。
JEP 的结构保持尽可能轻量级,以便开始,并且可能会在以后扩展。
何时应该使用 JEP#
当您的更改需要设计文档时。我们更倾向于将设计收集为 JEP,以便更好地发现和进一步参考。
当您的更改需要广泛的讨论时。在 issue 或 pull request 上进行相对简短的讨论是可以的,但是当讨论变得更长时,这对以后的消化是不切实际的。JEP 允许更新主文档,其中包含讨论的摘要,并且这些更新可以在添加 JEP 的 pull request 中进行讨论。
如何启动 JEP#
首先,创建一个带有 JEP 标签的 issue。所有与 JEP 相关的 pull request(即添加 JEP 本身以及任何实现 pull request)都应链接到此 issue。
然后创建一个 pull request,添加一个名为 %d-{短标题}.md 的文件 - 其中数字是 issue 的编号。
- 263: JAX PRNG 设计
- 2026: JAX 可转换函数的自定义 JVP/VJP 规则
- 4008: 自定义 VJP 和 `nondiff_argnums` 更新
- 4410: Omnistaging
- 9263: 类型化键和可插拔 RNG
- 9407: JAX 类型提升语义的设计
- 9419: Jax 和 Jaxlib 版本控制
- 10657: 在 JAX 中排序副作用
- 11830: `jax.remat` / `jax.checkpoint` 新实现
- 12049: JAX 的类型注释路线图
- 14273: 用于简单设备代码的 `shard_map` (`shmap`)
- 15856: `jax.extend`,一个扩展模块
- 17111: 高效转置 `shard_map` (和其他映射)
- 18137: JAX NumPy 和 SciPy 包装器的范围
- 25516: 基于工作量的版本控制
一些早期的 JEP 是从其他文档、issue 和 pull request 中事后转换而来的,因此它们可能不完全反映上述流程。