JAX 增强提案 (JEP)#
大多数更改都可以通过简单的 issue/discussion 和 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-{short-title}.md 的文件 — 数字应为 issue 编号。
- 263:JAX PRNG 设计
- 2026:JAX 可转换函数的自定义 JVP/VJP 规则
- 4008:自定义 VJP 和
nondiff_argnums更新 - 4410:全阶段
- 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:基于工作量的版本控制
- 28661:支持
__jax_array__协议
一些早期的 JEP 是在事后从其他文档、issue 和 pull request 转换而来的,因此它们可能不完全反映上述流程。