jax-cuda13-plugin
JAX Plugin for NVIDIA GPUs
Decision gist · record as of 2026-08-14
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.AI-flagged interpretation of the facts on this page — verify before relying
Before you install
- Requires NVIDIA CUDA 13 toolkit and compatible GPU; only available on Linux x86_64 and aarch64; requires Python 3.12 or later.
- 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 · maintenance · safety
Apache-2.0 (permissive) — Apache-2.0 permissive license allows commercial and private use with minimal restrictions; suitable for most projects.
last release 2026-07-16 (29 days) · last repo commit 2026-08-14 · 36,159 stars
0 known vulnerabilities (OSV.dev, 2026-08-14) · 420,823 downloads/mo, #6,787 on PyPI
Alternatives
Verify before relying
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)- 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
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 on this page. Verify before relying on it.
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
jax-cuda13-plugin on PyPI
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.
Requires NVIDIA CUDA 13 toolkit and compatible GPU; only available on Linux x86_64 and aarch64; requires Python 3.12 or later.
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)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 the current Python release >=3.12 |
| Install friction | Medium. Platform-specific wheel |
| Runtime dependencies | 1 packagejax-cuda13-pjrt |
| Maintenance | Actively maintained 29 days since the last release |
| Last repo commit | |
| First released | |
| Downloads | 420,823 / month, #6,787 on PyPI 30-day window, as of 2026-08-14 |
| Known vulnerabilities | None known OSV.dev, checked 2026-08-14 |
| Classifiers | Development Status :: 5 - Production/StableProgramming Language :: Python :: 3.12Programming Language :: Python :: 3.13Programming Language :: Python :: 3.14Programming Language :: Python :: Free Threading :: 3 - Stable |
Evidence: jax_cuda13_plugin-0.11.0-cp312-cp312-manylinux_2_27_aarch64.whl; jax_cuda13_plugin-0.11.0-cp312-cp312-manylinux_2_27_x86_64.whl; jax_cuda13_plugin-0.11.0-cp313-cp313-manylinux_2_27_aarch64.whl; jax_cuda13_plugin-0.11.0-cp313-cp313-manylinux_2_27_x86_64.whl; jax_cuda13_plugin-0.11.0-cp314-cp314-manylinux_2_27_aarch64.whl; jax_cuda13_plugin-0.11.0-cp314-cp314-manylinux_2_27_x86_64.whl; jax_cuda13_plugin-0.11.0-cp314-cp314t-manylinux_2_27_aarch64.whl; jax_cuda13_plugin-0.11.0-cp314-cp314t-manylinux_2_27_x86_64.whl
Tags
Let your AI agent find packages like this
Example. Real query, live index.
You found this page by searching. An agent finds it by wishing: SkillFed indexes 14,416 PyPI packages by what they can do, searchable in plain language.
wish › “jax nvidia gpu plugin”
- jax-cuda13-pluginProvides NVIDIA GPU support for JAX by enabling CUDA 13 compilation…
- jax-cuda12-pluginEnables JAX to run numerical computations and machine learning…
- jax-cuda12-pjrtProvides NVIDIA GPU acceleration for JAX numerical computing via the…
Give your agent the search over MCP, or paste the wish link into any chat.
More Artificial Intelligence packages
LiteLLM provides a unified Python interface to call 100+ LLM providers (OpenAI, Anthropic, Gemini, Bedrock, Azure, and others) using OpenAI-compatible API format, available as both a Python SDK and a self-hosted AI Gateway proxy server.
Install it if you need to work with multiple LLM providers or want to centralize LLM routing in your organization.
Client library and CLI tool for downloading, uploading, and managing models, datasets, and repositories on the Hugging Face Hub platform.
Install it if you work with Hugging Face Hub models or datasets.
LangChain provides a framework for building agents and LLM-powered applications by composing language models, tools, and memory through a unified API that abstracts over multiple model providers.
hf-xet provides chunk-based deduplication and efficient file transfer for the Hugging Face Hub, enabling faster uploads and downloads of large files with local disk caching.
Tokenizers converts raw text into token sequences for NLP models, with support for training custom vocabularies and using pre-built tokenizers (BPE, WordPiece) optimized for speed via Rust.
Transformers provides a unified framework for loading, fine-tuning, and running state-of-the-art pretrained models across text, vision, audio, video, and multimodal tasks using PyTorch, JAX, or TensorFlow.
Install it if you need to run or train any transformer-based model for NLP, vision, audio, or multimodal tasks.
See also jax-cuda12-plugin · jax-cuda12-pjrt · jax-cuda13-pjrt · tokamax · augmax · jaxlib · jax · klujax · jax-jumpy · blackjax