JAX 增强提案 (JEPs)#
大多数更改可以通过简单的议题/讨论和拉取请求进行讨论。
然而,有些更改的范围较大或需要更多讨论,这些更改应作为 JEP 实施。这允许编写更长的文档,这些文档本身可以在拉取请求中进行讨论。
JEP 的结构在开始时尽可能保持轻量级,并可能在以后进行扩展。
何时使用 JEP#
当您的更改需要设计文档时。我们倾向于将设计作为 JEP 收集,以提高可发现性并方便后续参考。
当您的更改需要广泛讨论时。在议题或拉取请求上进行相对简短的讨论是可以的,但当讨论变得更长时,这不利于后续的理解和消化。JEP 允许用讨论摘要更新主文档,并且这些更新本身可以在添加 JEP 的拉取请求中进行讨论。
如何启动 JEP#
首先,创建一个带有 JEP 标签的议题。所有与 JEP 相关的拉取请求(即添加 JEP 本身以及任何实现拉取请求)都应链接到此议题。
然后创建一个拉取请求,添加一个名为 %d-{short-title}.md 的文件——其中数字是议题编号。
- 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 是事后从其他文档、议题和拉取请求转换而来的,因此它们可能不完全反映上述过程。