The jax-cuda12-pjrt package provides the JAX XLA PJRT backend for NVIDIA GPUs, specifically built with CUDA 12. It serves as the `jaxlib` implementation when GPU acceleration is desired. The current version is 0.9.2, and JAX along with its ecosystem components typically follow a rapid release cadence, often with monthly or bi-monthly updates.
pip install jax jax-cuda12-pjrtVerified import paths — ran on the pinned version, not inferred.
This quickstart demonstrates basic JAX usage. It checks the JAX backend and available devices, then defines and executes a JIT-compiled function, confirming that the computation leverages GPU acceleration if available.
Refer to the JAX migration guide for `pmap` (docs.jax.dev/en/latest/migrate_pmap.html) and use `jax.shard_map` instead of `jax.pmap`.
Update code to pass an array-like object that implements the DLPack protocol, rather than a raw capsule, to `jax.dlpack.from_dlpack`.
Use `jax.lax.pcast(..., to='varying')` as the replacement for `jax.lax.pvary`.
To recover the previous host-based generation and ensure higher precision for narrow-width floats, explicitly cast the NumPy output: `jnp.array(np.arange(...))`.
Ensure that the `PartitionSpec` of your inputs precisely matches the `in_specs` argument when using `jax.shard_map` in 'Explicit' mode. Omit `in_specs` if you intend for `shard_map` to infer the partitioning.
Ensure only `jax` and your desired `jaxlib` variant (e.g., `jax-cuda12-pjrt` for CUDA 12 GPU or `jaxlib` for CPU) are installed. Uninstall any conflicting `jaxlib` packages before installing your target GPU backend.