jaxlib

What it is and what it does

jaxlib is the low-level compiled library that powers JAX's numerical computing engine. It wraps XLA (Accelerated Linear Algebra) to compile and execute array operations on accelerators—GPUs, TPUs, and CPUs—with automatic optimization and fusion. It depends on scipy, numpy, and ml_dtypes for numerical operations and type support.

jaxlib is essential infrastructure for anyone using JAX for machine learning, scientific computing, or large-scale numerical work. It handles the translation from Python code into compiled, device-optimized kernels. The package is actively maintained, supports Python 3.12, 3.13, and 3.14, and is available as prebuilt wheels for common platforms, making installation straightforward on supported systems.

Use it for:

• Running machine learning models on GPU or TPU clusters for training and inference at scale
• Accelerating numerical simulations and scientific computing workloads
• Enabling automatic differentiation and JIT compilation for high-performance gradient computation
• Deploying applications that require hardware acceleration on cloud or on-premise infrastructure

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

jaxlib is the compiled XLA backend that enables JAX to execute numerical computations on accelerators like GPUs, TPUs, and CPUs with high performance.

Yes, if you are using JAX. jaxlib is a required dependency for JAX to function; it cannot be avoided. It is actively maintained, has no known vulnerabilities, uses a permissive license, and offers prebuilt wheels for modern Python on common platforms. Install friction is moderate due to platform-specific compilation, but prebuilt binaries eliminate the need to compile from source on supported systems.

Install

jaxlib on PyPI

pip install jaxlib

uv add jaxlib

poetry add jaxlib

Installing jaxlib

Before you install

Medium install friction due to platform-specific compiled wheels; prebuilt binaries are available for Python 3.12, 3.13, and 3.14 across Linux x86_64, Linux aarch64, macOS arm64, and Windows x86_64. Active maintenance with recent releases.

License in practice

Apache-2.0 permissive license allows commercial and private use with minimal restrictions; you must include a copy of the license and state significant changes.

Quickstart

pip install jaxlib
import jaxlib
# jaxlib provides the compiled backend; typical usage is indirect
# through higher-level APIs that depend on scipy, numpy, ml_dtypes

Requires Python 3.12 or later; compiled wheels are platform-specific (Linux x86_64/aarch64, macOS arm64, Windows x86_64).

Verify before relying

• Whether jaxlib can be installed from source on platforms without prebuilt wheels
• GPU/TPU driver and runtime version requirements for hardware acceleration
• Performance characteristics and memory overhead relative to alternatives
• Specific CUDA, ROCm, or other accelerator framework version compatibility

Package facts

Evidence: jaxlib-0.11.0-cp312-cp312-macosx_11_0_arm64.whl; jaxlib-0.11.0-cp312-cp312-manylinux_2_27_aarch64.whl; jaxlib-0.11.0-cp312-cp312-manylinux_2_27_x86_64.whl; jaxlib-0.11.0-cp312-cp312-win_amd64.whl; jaxlib-0.11.0-cp313-cp313-macosx_11_0_arm64.whl; jaxlib-0.11.0-cp313-cp313-manylinux_2_27_aarch64.whl; jaxlib-0.11.0-cp313-cp313-manylinux_2_27_x86_64.whl; jaxlib-0.11.0-cp313-cp313-win_amd64.whl; jaxlib-0.11.0-cp314-cp314-macosx_11_0_arm64.whl; jaxlib-0.11.0-cp314-cp314-manylinux_2_27_aarch64.whl; jaxlib-0.11.0-cp314-cp314-manylinux_2_27_x86_64.whl; jaxlib-0.11.0-cp314-cp314t-macosx_11_0_arm64.whl; jaxlib-0.11.0-cp314-cp314t-manylinux_2_27_aarch64.whl; jaxlib-0.11.0-cp314-cp314t-manylinux_2_27_x86_64.whl; jaxlib-0.11.0-cp314-cp314-win_amd64.whl

Tags

More Scientific/Engineering packages