jaxlib

library0.11.1pypypi✓ verified 24d ago

jaxlib is the essential support library for JAX, containing the binary (C/C++) parts of the JAX ecosystem, including Python bindings, the XLA compiler, the PJRT runtime, and various handwritten kernels. While JAX itself is a pure Python package providing the high-level API, jaxlib acts as its compiled backend, enabling high-performance numerical computation on CPUs, GPUs, and TPUs. The current version is 0.9.2, and it follows a frequent release cadence, often aligning with or preceding JAX releases.

ai-ml data

pip install --upgrade pip
pip install --upgrade jax jaxlib

github ↗package ↗

INSTALL

IMPORT

SIG · JAXLIB

jaxlib

ai-mlpythonv0.11.1

Install

18.8s avg

Import

2016ms

Disk

587MB

Pass rate

5/ 10

Env Coverage5 / 10

glibc

3.9–3.13

musl

3.9–3.13

Install & Compatibility

Where this runs

tested against v0.6.2 · pip install

no network on importno background threads

Install × environment matrix

Each cell = how many times install + import succeeded across repeated harness runs. Partial = flaky.
glibc = Debian/Ubuntu slim · musl = Alpine Linux

musl

glibc

py 3.10

1/2 runs

✓ 12.55s

py 3.11

1/2 runs

✓ 13.9s

py 3.12

1/2 runs

✓ 14.05s

py 3.13

1/2 runs

✓ 14.4s

py 3.9

1/2 runs

✓ 38.85s

587MB installed

● package 587MB

Code

Verified usage

Verified import paths — ran on the pinned version, not inferred.

jax

✓ import jax

Users typically import `jax` (the Python frontend) to interact with the JAX ecosystem, which implicitly uses `jaxlib` as its backend. Direct imports from `jaxlib` for general user code are rare.

jax.numpy

✓ import jax.numpy as jnp

JAX's NumPy-like API (`jax.numpy`) is the primary way users interact with array operations that are compiled and executed by `jaxlib`.

This quickstart demonstrates a basic JAX program that implicitly leverages `jaxlib` for Just-In-Time (JIT) compilation and execution on available accelerators (CPU, GPU, or TPU). It defines a simple numerical function, compiles it with `jax.jit`, and performs an operation on a JAX array. The output shows detected devices and a sample of the computation.

import jax import jax.numpy as jnp def my_function(x): return jnp.sin(x) * jnp.cos(x) # JIT-compile the function for performance compiled_function = jax.jit(my_function) # Create a JAX array x = jnp.linspace(0, 10, 1000) # Run the compiled function y = compiled_function(x) print(f"JAX detected devices: {jax.devices()}") print(f"Result array shape: {y.shape}") print(f"First 5 elements of y: {y[:5]}")

Debug

Known issues

breakingThe `jax.pmap` function is now in maintenance mode, and its default implementation has changed. Users are strongly encouraged to migrate new code to `jax.shard_map` for data parallelism.

fix

Rewrite `jax.pmap` usage to `jax.shard_map`. Consult the JAX migration guide for `pmap`.

affects: >=0.8.0

breakingThe minimum supported NumPy version is now 2.0, and consequently, the minimum supported SciPy version is 1.13. Using older versions will lead to errors.

fix

Upgrade NumPy to version 2.0 or newer and SciPy to 1.13 or newer (`pip install --upgrade numpy scipy`).

affects: >=0.7.2

breaking`jax.dlpack.from_dlpack` no longer accepts a raw DLPack capsule directly. It must now be called with an array implementing the `__dlpack__` and `__dlpack_device__` protocols.

fix

Ensure that the input to `jax.dlpack.from_dlpack` is an object that implements the DLPack Python protocol.

affects: >=0.7.2

gotchaJAX arrays are immutable, unlike NumPy arrays. In-place modification operations common in NumPy (e.g., `arr[0] = 5`) are not supported and will raise an error or require explicit functional updates like `.at[idx].set(value)`.

fix

Adopt a functional programming style for array manipulations. Use `array = array.at[index].set(value)` for updates, which returns a new array with the modification.

affects: All versions

gotchaInstalling `jaxlib` for NVIDIA GPUs or TPUs requires specific installation commands (e.g., `jax[cudaXX_pip]`) and often requires matching CUDA/cuDNN versions. Using `pip install jaxlib` alone will typically install a CPU-only version, and mismatched versions can lead to runtime errors or devices not being detected.

fix

Always refer to the official JAX installation guide for the correct, platform-specific commands. Ensure your CUDA/cuDNN versions are compatible with the `jaxlib` wheel you are installing.

affects: All versions

deprecatedThe semi-private type `jax._src.literals.TypedNdArray` is now a subclass of `np.ndarray`, not just a duck type. This change may affect code relying on `isinstance(x, np.ndarray)` or similar type checks for JAX internal types if they were previously treated as distinct from `np.ndarray`.

fix

If your code relies on distinguishing `TypedNdArray` from `np.ndarray`, adjust type checks accordingly. Convert to classic NumPy arrays using `np.asarray(x)` if necessary.

affects: >=0.9.2

Upgrade

Version history

0.11.1latest on PyPI · released Aug 17, 2026

Audit

Dependencies

jaxrequiredjaxlib is the compiled backend for the JAX Python frontend library; they are interdependent.

numpyrequiredRequired for array operations; minimum version 2.0 is now enforced for recent jaxlib versions.

scipyrequiredRequired for scientific computing functions; minimum version 1.13 is now enforced for recent jaxlib versions due to NumPy 2.0.

ml_dtypesrequiredRequired for machine learning specific dtypes.

Agent activity

12 hits · last 30 days

node

Resources

githubhttps://github.com/jax-ml/jax ↗

packagehttps://pypi.org/project/jaxlib/ ↗