jaxlib

jaxlib is the native support library underneath JAX: while JAX itself is a pure-Python package of composable function transformations (grad, jit, vmap, pmap), jaxlib supplies the compiled C++ parts those transformations ultimately run on — Python bindings to the XLA compiler and PJRT runtime, plus a set of handwritten CPU and GPU kernels for operations like linear algebra, sparse solves, and RNNs. It is built out of the same jax-ml/jax monorepo but shipped as its own PyPI wheel because it needs a native (per-platform, per-accelerator) build via Bazel, whereas jax is pure Python and installs anywhere.

Most users never import jaxlib directly; it’s installed as a version-pinned dependency of jax and loaded implicitly the moment jax.jit or jax.numpy is used. Its versioning is deliberately decoupled from jax’s own release cadence (jax pins a minimum jaxlib version) so that the pure-Python API and the native runtime can evolve — and be validated across CPU/CUDA/ROCm/TPU backends in CI — somewhat independently.

What You Get

• Python bindings to XLA and the PJRT client/executable/device abstractions that jax.jit compiles against
• Hand-written native kernels for CPU linear algebra (LAPACK-backed), sparse solves, and tridiagonal solves
• GPU kernel bindings (BLAS, solver, sparse, RNN, Triton) shared across CUDA and ROCm backends via a common plugin interface
• A PJRT plugin architecture (jax_plugins/) letting a single jax install target CPU, GPU, or TPU by swapping the installed jaxlib/plugin wheel
• DLPack interoperability and sharding/device-list primitives used by jax’s distributed and multi-device APIs

Common Use Cases

• Installed automatically as a pinned dependency whenever pip install jax[cuda] or similar variants are used
• Selecting a specific accelerator backend (CPU-only, CUDA, ROCm) by choosing the matching jaxlib/plugin wheel
• Debugging low-level compiled-execution or device-placement issues that surface below the jax Python API
• Building or packaging custom JAX distributions (e.g. for a specific CUDA version or TPU runtime) via the Bazel build

Under The Hood

Architecture jaxlib is organized as a set of C++ extension modules (bound to Python via nanobind, per headers like nb_class_ptr.h and py_array.h) that expose the XLA/PJRT client, executable, and array/device abstractions (py_executable.h, py_array.h, py_device_list.h, sharding.h) consumed by the pure-Python jax package. Hardware-specific code is split by backend: jaxlib/cpu holds LAPACK-based linear-algebra and sparse kernels, while jaxlib/gpu and jaxlib/cuda hold shared GPU kernel wrappers (BLAS, solver, sparse, RNN, Triton) built against CUDA or ROCm. Bazel (MODULE.bazel, per-directory BUILD files, jax_python_wheel.bzl) drives the entire build, and jax_plugins/ implements a PJRT plugin interface so a single jax checkout can target CPU, GPU, or TPU simply by swapping which jaxlib/plugin wheel is installed — meaning changes to the core client/executable abstractions ripple through both this C++ layer and the Python package that wraps it.

Tech Stack The library is written primarily in C++ bound to Python via nanobind, compiled with Bazel rather than a standard Python build backend, with a thin setuptools/wheel layer (pyproject.toml, setup.py, build_wheel.py) packaging the Bazel-built shared objects into per-platform, per-accelerator pip wheels. Core native dependencies are XLA and the PJRT runtime (vendored under third_party/), plus optional CUDA/ROCm/oneAPI toolkits for GPU kernels and LAPACK for CPU linear algebra. CI (.github/workflows) runs dozens of matrix configurations spanning CPU, CUDA (including H100/B200), ROCm, oneAPI, and TPU, reflecting the native, multi-backend nature of the build.

Code Quality Testing is distributed across roughly 296 test files in the monorepo, with jaxlib-specific tests such as config_test.py, exception_test.py, and weakref_lru_cache_test.py sitting alongside the C++ modules; most correctness coverage runs through Bazel-driven CI (bazel_cpu.yml, bazel_cuda.yml, bazel_test_tpu.yml) and pytest workflows (pytest_cpu.yml, pytest_cuda.yml, pytest_tpu.yml) rather than a single unified test command, given the per-backend native build. Linting is enforced through a pre-commit configuration and a dedicated lint workflow, with ruff and pyrefly static typing configured in pyproject.toml covering the Python-facing surface; the C++ layer follows status-based (rather than exception-based) error handling typical of XLA/absl-style codebases.

API Design jaxlib’s distinguishing choice is exposing a narrow, stable PJRT plugin ABI so the Python-facing jax package can target CPU, CUDA, ROCm, or TPU by installing a different jaxlib/plugin wheel without any change to jax’s own code — a plugin model most numerical Python libraries don’t attempt. Its public surface is intentionally minimal (device, array, executable, and sharding primitives, plus DLPack interop), since it exists to be wrapped by jax’s documented Python API rather than used directly, so day-to-day developer ergonomics live in the parent jax package rather than jaxlib’s own docs.