Registry / ai-ml / jax-cuda13-pjrt

jax-cuda13-pjrt

library0.10.0pypypi✓ verified 79d ago

JAX PJRT plugin for NVIDIA GPUs using CUDA 13. Provides the XLA compiler backend for JAX on NVIDIA hardware. Version 0.10.0 aligns with JAX v0.10.0. Release cadence follows JAX releases.

ai-ml

pip install jax-cuda13-pjrt

github ↗package ↗

INSTALL

IMPORT

SIG · JAX-CUDA13-PJRT

jax-cuda13-pjrt

ai-mlpythonv0.10.0

harness data pending

Install & Compatibility

Where this runs

No compatibility data collected yet for this library.

Code

Verified usage

Verified import paths — ran on the pinned version, not inferred.

jax_plugins

✓ import jax_plugins

✗ import jax_cuda13_pjrt

Import the plugin to register the CUDA 13 backend, then use JAX as usual.

import jax import jax.numpy as jnp import jax_cuda13_pjrt # noqa: F401, ensure NVIDIA GPU plugin is registered # Verify device is visible print(jax.devices()) # Simple computation x = jnp.array([1, 2, 3]) y = jnp.square(x) print(y)

Debug

Known issues

gotchaThe plugin must be imported before any JAX computation to ensure the GPU backend is selected.

fix

Always import jax_cuda13_pjrt at the top of your script.

affects: all

breakingPlugin version must match the JAX version. Using mismatched versions may cause runtime errors or undefined behavior.

fix

Ensure jax==0.10.0 and jax-cuda13-pjrt==0.10.0 are installed together.

affects: all

gotchaThe plugin is only for CUDA 13.x; it will not work with older CUDA toolkits.

fix

Verify your CUDA version with nvcc --version. If CUDA < 13, use jax-cuda12-pjrt instead.

affects: all

Upgrade

Version history

0.10.0latest on PyPI · released Apr 16, 2026

Audit

Dependencies

jaxrequiredJAX core is required; the plugin is a backend.

jaxlibrequiredRequired for XLA runtime; must match CUDA version.

Agent activity

6 hits · last 30 days

node

Resources

githubhttps://github.com/jax-ml/jax ↗

packagehttps://pypi.org/project/jax-cuda13-pjrt/ ↗