---
id: jax-cuda13-pjrt
version: "0.11.0"
license: Apache-2.0
license_treatment: permissive
maintenance: active
---

# jax-cuda13-pjrt — JAX XLA PJRT Plugin for NVIDIA GPUs

License: permissive · Maintenance: active · Downloads: 415.2K/mo

## 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 above — 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

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

- License: Apache-2.0 (permissive)
- Python support: unspecified
- Install friction: medium
- Maintenance: active
- Downloads: 415.2K/month (top 15,000 on PyPI)
- Known vulnerabilities: none known

## Tags

JAX NVIDIA GPU acceleration, PJRT CUDA plugin, JAX XLA GPU compilation, CUDA 13 JAX support, GPU-accelerated array computing, gpu-acceleration, machine-learning, numerical-computing

[View on SkillFed](https://skillfed.io/packages/jax-cuda13-pjrt) · [View on PyPI](https://pypi.org/project/jax-cuda13-pjrt/)