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

# jax-cuda13-plugin — JAX Plugin for NVIDIA GPUs

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

## What it is and what it does

jax-cuda13-plugin is a platform-specific plugin that enables JAX to compile and execute numerical computations on NVIDIA GPUs using CUDA 13. It acts as a bridge between JAX's compiler infrastructure and NVIDIA's GPU hardware, allowing automatic differentiation, JIT compilation, and vectorization operations to run on GPUs for significant performance gains in machine learning and scientific computing workloads.

The plugin is tightly integrated with JAX's ecosystem and depends on jax-cuda13-pjrt for the underlying PJRT implementation. It is distributed as pre-compiled wheels for Linux systems (both x86_64 and aarch64 architectures) and supports Python 3.12, 3.13, and 3.14, including free-threaded variants. Installation requires CUDA 13 to be present on the system.

Use it for:

- Accelerate machine learning training on NVIDIA GPUs by enabling JAX's JIT compilation and automatic differentiation to target GPU hardware.
- Run numerical simulations and scientific computations that benefit from GPU parallelization through JAX's vmap and pmap transformations.
- Deploy JAX-based inference pipelines on GPU clusters for production machine learning systems.
- Develop and test GPU-accelerated JAX code on Linux systems with NVIDIA GPUs before scaling to multi-GPU environments.

## Worth the install?

AI-flagged interpretation of the facts above — verify before relying.

Provides NVIDIA GPU support for JAX by enabling CUDA 13 compilation and execution on compatible GPUs through the jax-cuda13-pjrt runtime.

Yes, if you are using JAX on Linux with an NVIDIA GPU and CUDA 13 installed. This plugin is essential for GPU acceleration and is actively maintained with no known vulnerabilities. Install friction is moderate but manageable for systems meeting the platform requirements. Not applicable on non-Linux systems or without compatible NVIDIA hardware.

## Install

pip install jax-cuda13-plugin
uv add jax-cuda13-plugin
poetry add jax-cuda13-plugin

## Installing jax-cuda13-plugin

Before you install: Medium install friction due to platform-specific wheels (x86_64 and aarch64 Linux only) and a single runtime dependency on jax-cuda13-pjrt. Actively maintained with recent releases and no known vulnerabilities.

License in practice: Apache-2.0 permissive license allows commercial and private use with minimal restrictions; suitable for most projects.

Quickstart:

pip install jax-cuda13-plugin

import jax
import jax.numpy as jnp

# GPU acceleration is automatically available
x = jnp.ones((100, 100))
result = jax.jit(lambda a: a @ a)(x)

Requires NVIDIA CUDA 13 toolkit and compatible GPU; only available on Linux x86_64 and aarch64; requires Python 3.12 or later.

Verify before relying:

- Whether this plugin is required separately or included in main JAX installations for GPU users
- Specific NVIDIA GPU models and driver versions supported by CUDA 13
- Performance characteristics compared to alternative JAX GPU backends

## Package facts

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

## Tags

jax nvidia gpu plugin, cuda 13 jax acceleration, jax gpu support, nvidia cuda jax runtime, jax pjrt cuda plugin, gpu-accelerated jax, jax cuda13 backend, gpu-acceleration, cuda-plugin, machine-learning

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