jax-cuda13-pjrt

What it is and what it does

This package is a PJRT (Portable Runtime for JAX) plugin that bridges JAX's numerical computing framework to NVIDIA GPUs via CUDA 13. It enables JAX's core transformations—automatic differentiation (grad), JIT compilation, and vectorization (vmap)—to execute on NVIDIA hardware by interfacing with XLA's compiler infrastructure.

When installed, it allows JAX programs to transparently offload array computations and neural network training to NVIDIA GPUs without code changes. The plugin handles device memory management, kernel dispatch, and multi-device coordination. It is a platform-specific binary distribution (wheels for Linux x86_64 and aarch64 only) that must match your system's CUDA 13 installation.

Use it for:

• Train neural networks on NVIDIA GPUs using JAX's automatic differentiation and JIT compilation.
• Accelerate scientific computing workloads (linear algebra, FFTs, stochastic simulations) on GPU hardware.
• Run JAX-based machine learning research code on NVIDIA clusters or cloud instances with CUDA 13.
• Develop high-performance numerical applications requiring GPU parallelism using vmap.

Worth the install?

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

Provides NVIDIA GPU acceleration for JAX numerical computing via the PJRT plugin interface, enabling XLA compilation and execution on CUDA 13 hardware.

Yes, if you have NVIDIA GPU hardware with CUDA 13 and want to use JAX for accelerated computing. The package is actively maintained, permissively licensed, and has no known vulnerabilities. Install friction is moderate due to platform specificity and CUDA system dependency, but necessary for GPU support. Not applicable on CPU-only or non-NVIDIA systems.

Install

jax-cuda13-pjrt on PyPI

pip install jax-cuda13-pjrt

uv add jax-cuda13-pjrt

poetry add jax-cuda13-pjrt

Installing jax-cuda13-pjrt

Before you install

Medium install friction due to platform-specific wheels (manylinux_2_27 x86_64 and aarch64 only) and CUDA 13 system dependency. Active maintenance with recent release (29 days old) and strong upstream project health (36159 stars).

License in practice

Apache-2.0 permissive license allows commercial and private use with minimal restrictions; suitable for most production and research contexts.

Quickstart

pip install jax-cuda13-pjrt
import jax
import jax.numpy as jnp

# JAX operations now dispatch to NVIDIA GPU via PJRT
x = jnp.ones((100, 100))
y = jnp.dot(x, x)  # runs on GPU

Requires CUDA 13 runtime libraries installed on the system and an NVIDIA GPU; only available on Linux x86_64 or aarch64 platforms.

Verify before relying

• Whether this package is required separately or automatically installed as a JAX dependency on compatible systems
• Specific NVIDIA GPU models and compute capabilities supported by CUDA 13 PJRT plugin
• Performance characteristics or overhead compared to other JAX GPU backends

Package facts

Evidence: jax_cuda13_pjrt-0.11.0-py3-none-manylinux_2_27_aarch64.whl; jax_cuda13_pjrt-0.11.0-py3-none-manylinux_2_27_x86_64.whl

Tags

More Scientific/Engineering packages

Further reading