API compatibility
Contents
API compatibility#
JAX is constantly evolving, and we want to be able to make improvements to itsAPIs. That said, we want to minimize churn for the JAX user community, and wetry to make breaking changes rarely.
JAX Versioning#
JAX usesEffort-based versioning (seeJEP 25516: Effort-based versioning for JAX), and is currently in the Zero version phase.This means that for version0.X.Y, incrementingY will introduce minorbreaking changes, and incrementingX will introduce major breaking changes.
For any breaking change, JAX currently follows a 3 month deprecation policy.When an incompatible change is made to an API, we will make our best effortto obey the following procedure:
the change will be announced in
CHANGELOG.mdand in the doc string for thedeprecated API, and the old API will issue aDeprecationWarning.three months after the
jaxrelease that deprecated an API, we may remove thedeprecated API at any time. Note that three months is alower bound, and isintentionally chosen to be faster than that of many more mature projects. Inpractice, deprecations may take considerably longer, particularly if there aremany users of a feature. If a three month deprecation period becomesproblematic, please raise this with us.
We reserve the right to change this policy at any time.
What is covered?#
Only public JAX APIs are covered, which includes the following modules:
jaxjax.dlpackjax.imagejax.laxjax.nnjax.numpyjax.opsjax.profilerjax.random(seedetails below)jax.scipyjax.treejax.tree_utiljax.test_util
Not everything in these modules is intended to be public, and over time, weare working to separate public and private APIs. Public APIs are documentedin the JAX documentation.Additionally, our goal is that all non-public APIs should have namesprefixed with underscores, although we do not entirely comply with this yet.
What is not covered?#
Explicitly private APIs#
Any API or import path prefixed with an underscore is explicitly private,and may change without warning between JAX releases. We are working to moveall private APIs intojax._src to make these expectations more clear.
jaxlib#
Any import path in thejaxlib package is considered private, and may changewithout warning between releases. Some APIs defined injaxlib have publicaliases in thejax package.
Legacy internal APIs#
In addition, there are several legacy modules that currently expose someprivate APIs without an underscore, including:
jax.corejax.interpretersjax.libjax.util
We are actively working on deprecating these modules and the APIs they contain.In most cases, such deprecations will follow the 3 month deprecation period,but this may not always be possible. If you use any such APIs, please expectthem to be deprecated soon, and seek alternatives.
Experimental and example libraries#
The following modules include code for experimental or demonstration purposes,and API may change between releases without warning:
jax.experimentaljax.example_libraries
We understand that some users depend onjax.experimental, and so in most caseswe follow the 3 month deprecation period for changes, but this may not always bepossible.
JAX extend#
Thejax.extend module includes semi-public JAX internal APIs that aremeant for use by downstream projects, but do not have the same stabilityguarantees of the main JAX package. If you have code that usesjax.extend,we would strongly recommend CI tests against JAX’s nightly releases, so as tocatch potential changes before they are released.
For details onjax.extend, see thejax.extend module documentation, or the design document,jax.extend: a module for extensions.
Numerics and randomness#
Theexact values of numerical operations are not guaranteed to bestable across JAX releases. In fact, exact numerics are notnecessarily stable at a given JAX version, across acceleratorplatforms, within or withoutjax.jit, and more.
For a fixed PRNG key input, the outputs of pseudorandom functions injax.random may vary across JAX versions. The compatibility policyapplies only to the outputdistribution. For example, the expressionjax.random.gumbel(jax.random.key(72)) may return a different valueacross JAX releases, butjax.random.gumbel will remain apseudorandom generator for the Gumbel distribution.
We try to make such changes to pseudorandom values infrequently. Whenthey happen, the changes are announced in the changelog, but do notfollow a deprecation cycle. In some situations, JAX might expose atransient configuration flag that reverts the new behavior, to helpusers diagnose and update affected code. Such flags will last adeprecation window’s amount of time.