JAX Enhancement Proposals (JEPs)
Contents
JAX Enhancement Proposals (JEPs)#
Most changes can be discussed with simple issues/discussions and pull requests.
Some changes though are a bit larger in scope or require more discussion, andthese should be implemented as JEP. This allows for writing longer documentsthat can be discussed in a pull request themselves.
The structure of JEPs is kept as lightweight as possible to start and mightbe extended later on.
When you should use a JEP#
When your change requires a design doc. We prefer collecting the designs asJEPs for better discoverability and further reference.
When your change requires extensive discussion. It’s fine to have relativelyshort discussions on issues or pull requests, but when the discussion getslonger this becomes unpractical for later digestion. JEPs allow to update themain document with a summary of the discussion and these updates can bediscussed themselves in the pull request adding the JEP.
How to start a JEP#
First, create an issue with theJEP label. All pull requests that relate tothe JEP (i.e. adding the JEP itself as well as any implementing pull requests)should be linked to this issue.
Then create a pull request that adds a file named%d-{short-title}.md - with the number being the issue number.
- 263: JAX PRNG Design
- 2026: Custom JVP/VJP rules for JAX-transformable functions
- 4008: Custom VJP and `nondiff_argnums` update
- 4410: Omnistaging
- 9263: Typed keys & pluggable RNGs
- 9407: Design of Type Promotion Semantics for JAX
- 9419: Jax and Jaxlib versioning
- 10657: Sequencing side-effects in JAX
- 11830: `jax.remat` / `jax.checkpoint` new implementation
- 12049: Type Annotation Roadmap for JAX
- 14273: `shard_map` (`shmap`) for simple per-device code
- 15856: `jax.extend`, an extensions module
- 17111: Efficient transposition of `shard_map` (and other maps)
- 18137: Scope of JAX NumPy & SciPy Wrappers
- 25516: Effort-based versioning
- 28661: Supporting the `__jax_array__` protocol
- 28845: Stateful Randomness in JAX
Several early JEPs were converted in hindsight from other documentation,issues, and pull requests, so they might not exactly reflect the processoutlined above.
