jax-cuda12-pjrt
JAX XLA PJRT Plugin for NVIDIA GPUs
Decision gist · record as of 2026-08-14
Yes, if you are using JAX on an NVIDIA GPU with CUDA 12 and need GPU acceleration. The plugin is actively maintained, carries no known vulnerabilities, and is essential for GPU-backed JAX workloads. Install friction is moderate due to platform specificity and CUDA 12 runtime dependency, but this is expected for GPU-accelerated packages. Not needed if you are using CPU-only JAX or a different GPU backend.AI-flagged interpretation of the facts on this page — verify before relying
Before you install
- Requires NVIDIA GPU with CUDA 12 support and CUDA 12 runtime libraries installed on the system; only available for Linux x86_64 and aarch64 platforms.
- Medium install friction due to platform-specific wheels (manylinux_2_27 x86_64 and aarch64 only) and CUDA 12 runtime dependency.
- Active maintenance with recent releases; 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 production and research deployments.
last release 2026-07-16 (29 days) · last repo commit 2026-08-14 · 36,158 stars
0 known vulnerabilities (OSV.dev, 2026-08-14) · 774,416 downloads/mo, #5,095 on PyPI
Alternatives
Verify before relying
pip install jax-cuda12-pjrt
import jax
import jax.numpy as jnp
# JAX operations now run on NVIDIA GPU
result = jax.jit(lambda x: jnp.dot(x, x))(jnp.ones((100, 100)))- Whether this package is a standalone plugin or requires a separate JAX installation as a prerequisite.
- Specific NVIDIA GPU compute capability requirements beyond CUDA 12 support.
- Performance characteristics compared to other JAX GPU backends or configurations.
What it is and what it does
jax-cuda12-pjrt is a plugin that connects JAX's XLA compiler to NVIDIA GPUs via the PJRT (Portable Runtime) interface, enabling GPU-accelerated execution of JAX programs on CUDA 12 hardware. It allows you to write numerical code using JAX's NumPy-like API and automatic differentiation, then compile and run it efficiently on NVIDIA accelerators.
The package is a specialized distribution component of JAX itself—it bundles the PJRT runtime plugin for CUDA 12 in a platform-specific wheel. Installation brings GPU support to JAX's core transformations (grad, jit, vmap) without requiring manual compilation of the plugin. It is actively maintained and carries no known security vulnerabilities.
Use it for
- Training deep learning models on NVIDIA GPUs using JAX's automatic differentiation and JIT compilation.
- Running large-scale numerical simulations and scientific computing workloads with GPU acceleration.
- Developing machine learning research code that needs efficient gradient computation across GPU clusters.
- Accelerating array operations and linear algebra computations on NVIDIA hardware without writing CUDA code.
- Deploying JAX-based inference pipelines that require GPU performance for production inference.
Worth the install?
AI-flagged interpretation of the facts on this page. Verify before relying on it.
Yes, if you are using JAX on an NVIDIA GPU with CUDA 12 and need GPU acceleration.
The plugin is actively maintained, carries no known vulnerabilities, and is essential for GPU-backed JAX workloads. Install friction is moderate due to platform specificity and CUDA 12 runtime dependency, but this is expected for GPU-accelerated packages. Not needed if you are using CPU-only JAX or a different GPU backend.
Install
jax-cuda12-pjrt on PyPI
Before you install
Medium install friction due to platform-specific wheels (manylinux_2_27 x86_64 and aarch64 only) and CUDA 12 runtime dependency. Active maintenance with recent releases; no known vulnerabilities.
Requires NVIDIA GPU with CUDA 12 support and CUDA 12 runtime libraries installed on the system; only available for Linux x86_64 and aarch64 platforms.
License in practice
Apache-2.0 permissive license allows commercial and private use with minimal restrictions; suitable for most production and research deployments.
Quickstart
pip install jax-cuda12-pjrt
import jax
import jax.numpy as jnp
# JAX operations now run on NVIDIA GPU
result = jax.jit(lambda x: jnp.dot(x, x))(jnp.ones((100, 100)))Verify before relying
- Whether this package is a standalone plugin or requires a separate JAX installation as a prerequisite.
- Specific NVIDIA GPU compute capability requirements beyond CUDA 12 support.
- Performance characteristics compared to other JAX GPU backends or configurations.
Package facts
| License | Apache-2.0 permissive |
| Python support | Not specified |
| Install friction | Medium. Platform-specific wheel |
| Runtime dependencies | None |
| Maintenance | Actively maintained 29 days since the last release |
| Last repo commit | |
| First released | |
| Downloads | 774,416 / month, #5,095 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 :: 3Programming Language :: Python :: Free Threading :: 3 - Stable |
Evidence: jax_cuda12_pjrt-0.11.0-py3-none-manylinux_2_27_aarch64.whl; jax_cuda12_pjrt-0.11.0-py3-none-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 › “pjrt cuda 12 plugin”
- jax-cuda12-pjrtProvides NVIDIA GPU acceleration for JAX numerical computing via the…
- jax-cuda13-pjrtProvides NVIDIA GPU acceleration for JAX numerical computing via the…
- jax-cuda13-pluginProvides NVIDIA GPU support for JAX by enabling CUDA 13 compilation…
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-cuda13-pjrt · tokamax · augmax · jax-cuda13-plugin · jax-cuda12-plugin · jaxlib · torchax · jmp · jax · jax-jumpy