Jax and Jaxlib versioning
Jax and Jaxlib versioning#
jaxlib separate packages?#
We publish JAX as two separate Python wheels, namely
jax, which is a pure
Python wheel, and
jaxlib, which is a mostly-C++ wheel that contains libraries
pieces of LLVM used by XLA,
MLIR infrastructure, such as the MHLO Python bindings.
JAX-specific C++ libraries for fast JIT and PyTree manipulation.
We distribute separate
jaxlib packages because it makes it easy to
work on the Python parts of JAX without having to build C++ code or even having
a C++ toolchain installed.
jaxlib is a large library that is not easy for
many users to build, but most changes to JAX only touch Python code. By
allowing the Python pieces to be updated independently of the C++ pieces, we
improve the development velocity for Python changes.
jaxlib is not cheap to build, but we want to be able to iterate on
and run the JAX tests in environments without a lot of CPU, for example in
Github Actions or on a laptop. Many of our CI builds simply use a prebuilt
jaxlib, rather than needing to rebuild the C++ pieces of JAX on each PR.
As we will see, distributing
jaxlib separately comes with a cost, in
that it requires that changes to
jaxlib maintain a backward compatible API.
However, we believe that on balance it is preferable to make Python changes
easy, even if at the cost of making C++ changes slightly harder.
jax’s version must be greater than or equal to
jaxlib’s version must be greater than or equal to the minimum
version specified by
jaxlib releases are numbered
x is the major
y is the minor version, and
z is an optional patch release.
Version numbers must follow
PEP 440. Version number comparisons
are lexicographic comparisons on tuples of integers.
jax release has an associated minimum
jaxlib version for
x.y.z must be no greater than
lx.ly.lz to be compatible,
the following must hold:
The jaxlib version (
lx.ly.lz) must be greater than or equal to the minimum jaxlib version (
The jax version (
x.y.z) must be greater than or equal to the jaxlib version (
These constraints imply the following rules for releases:
jaxmay be released on its own at any time, without updating
If a new
jaxlibis released, a
jaxrelease whose version is equal to or greater than version the
jaxlib’s version number must be made at the same time.
are currently checked by
jax at import time, instead of being expressed as
Python package version constraints.
jax checks the
jaxlib version at
runtime rather than using a
pip package version constraint because we
for a variety of hardware and software versions (e.g, GPU, TPU, etc.). Since we
do not know which is the right choice for any given user, we do not want
to install a
jaxlib package for us automatically.
In the future, we hope to separate out the hardware-specific pieces of
into separate plugins, at which point the minimum version could be expressed as
a Python package dependency. For now, we do provide
platform-specific extra requirements that install a compatible jaxlib version,
How can I safely make changes to the API of
jaxmay drop compatibility with older
jaxlibreleases at any time, so long as the minimum
jaxlibversion is increased to a compatible version. However, note that the minimum
jaxlib, even for unreleased versions of
jax, must be a released version! This allows us to use released
jaxlibwheels in our CI builds, and allows Python developers to work on
jaxat HEAD without ever needing to build
For example, to remove an old backwards compatibility path in the
jaxPython code, it is sufficient to bump the minimum jaxlib version and then delete the compatibility path.
jaxlibmay drop compatibility with older
jaxreleases lower than its own release version number. The version constraints enforced by
jaxwould forbid the use of an incompatible
For example, for
jaxlibto drop a Python binding API used by an older
jaxlibminor or major version number must be incremented.
If possible, changes to the
jaxlibshould be made in a backwards-compatible way.
jaxlibmay freely change its API, so long as the rules about
jaxbeing compatible with all
jaxlibs at least as new as the minimum version are followed. This implies that
jaxmust always be compatible with at least two versions of
jaxlib, namely, the last release, and the tip-of-tree version, effectively the next release. This is easier to do if compatibility is maintained, although incompatible changes can be made using version tests from
jax; see below.
For example, it is usually safe to add a new function to
jaxlib, but unsafe to remove an existing function or to change its signature if current
jaxis still using it. Changes to
jaxmust work or degrade gracefully for all
jaxlibreleases greater than the minimum up to HEAD.
Note that the compatibility rules here only apply to released versions of
jaxlib. They do not apply to unreleased versions; that is, it is ok
to introduce and then remove an API from
jaxlib if it is never released, or if
jax version uses that API.
How is the source to
jaxlib laid out?#
jaxlib is split across two main repositories, namely the
jaxlib/ subdirectory in the main JAX repository
and in the
XLA source tree, which lives inside the TensorFlow repository.
The JAX-specific pieces inside XLA are primarily in the
The reason that C++ pieces of JAX, such as Python bindings and runtime components, are inside the XLA tree is partially historical and partially technical.
The historical reason is that originally the
xla/python bindings were envisaged as general purpose Python bindings that
might be shared with other frameworks. In practice this is increasingly less
xla/python incorporates a number of JAX-specific pieces and is
likely to incorporate more. So it is probably best to simply think of
xla/python as part of JAX.
The technical reason is that the XLA C++ API is not stable. By keeping the
XLA:Python bindings in the XLA tree, their C++ implementation can be updated
atomically with the C++ API of XLA. It is easier to maintain backward and forward
compatibility of Python APIs than C++ ones, so
xla/python exposes Python APIs
and is responsible for maintaining backward compatibility at the Python
jaxlib is built using Bazel out of the
jax repository. The pieces of
jaxlib from the XLA repository are incorporated into the build
as a Bazel submodule.
To update the version of XLA used during the build, one must update the pinned
version in the Bazel
WORKSPACE. This is done manually on an
as-needed basis, but can be overriden on a build-by-build basis.
How do we make changes across the
jaxlib boundary between releases?#
The jaxlib version is a coarse instrument: it only lets us reason about releases.
However, since the
jaxlib code is split across repositories that
cannot be updated atomically in a single change, we need to manage compatibility
at a finer granularity than our release cycle. To manage fine-grained
compatibility, we have additional versioning that is independent of the
release version numbers.
We maintain an additional version number (
xla_client.py in the XLA repository.
The idea is that this version number, is defined in
together with the C++ parts of JAX, is also accessible to JAX Python as
jax._src.lib.xla_extension_version, and must
be incremented every time that a change is made to the XLA/Python code that has
backwards compatibility implications for
jax. The JAX Python code can then use
this version number to maintain backwards compatibility, e.g.:
# 123 is the new version number for _version in xla_client.py if jax._src.lib.xla_extension_version >= 123: # Use new code path ... else: # Use old code path.
Note that this version number is in addition to the constraints on the released version numbers, that is, this version number exists to help manage compatibility during development for unreleased code. Releases must also follow the compatibility rules given above.