Notes

This section contains shorter notes on topics relevant to using JAX; see also the longer design discussions in JAX Enhancement Proposals (JEPs).

Dependencies and version compatibility:

• API compatibility outlines JAX’s policies with regard to API compatibility across releases.

• Python and NumPy version support policy outlines JAX’s policies with regard to compatibility with Python and NumPy.

Memory and computation usage:

• Asynchronous dispatch describes JAX’s asynchronous dispatch model.

• Concurrency describes how JAX interacts with other Python concurrency.

• GPU memory allocation describes how JAX interacts with memory allocation on GPU.

Programmer guardrails:

• Rank promotion warning describes how to configure jax.numpy to avoid implicit rank promotion.

Arrays and data types:

• Type promotion semantics describes JAX’s implicit type promotion for functions of two or more values.

• Default dtypes and the X64 flag describes how JAX determines the default dtype for array creation functions.