JAX Plugin for NVIDIA GPUs (CUDA 12)

0.9.2 · active · verified Sun Apr 12

JAX is a Python library by Google for high-performance numerical computing, providing a NumPy-like interface with automatic differentiation and function transformations, capable of running on CPUs, GPUs, and TPUs. The `jax-cuda12-plugin` specifically provides NVIDIA GPU support for JAX, compatible with CUDA 12.x environments. JAX and its core library `jaxlib` (which this plugin extends) are actively maintained with frequent releases, typically on a monthly or bi-monthly schedule for minor versions.

Warnings

breaking JAX's parallel map (`jax.pmap`) is in maintenance mode and its default implementation has changed. New code is strongly encouraged to use `jax.shard_map` or `jax.jit` for automatic parallelism. The `auto=` parameter of `jax.experimental.shard_map.shard_map` was removed in v0.8.0.
Fix: Migrate from `jax.pmap` to `jax.shard_map` or `jax.jit`. Consult the JAX migration guide for details on `pmap` to `shard_map` transitions.
breaking Support for monolithic CUDA `jaxlibs` (e.g., `jaxlib==0.4.29+cuda12`) has been dropped. All CUDA support is now provided via plugin-based installations.
Fix: Ensure you are using the plugin-based installation, typically by installing `jax-cuda12-plugin` (or `jax[cuda12]`) as per the official JAX installation instructions.
gotcha JAX arrays are immutable. Attempting in-place modification (e.g., `x[0] = 10`) will result in a `TypeError`.
Fix: Use the `.at[]` syntax for element-wise updates, which returns a new array with the changes (e.g., `x = x.at[0].set(10)`).
gotcha `jax.numpy.arange` with a specified `step` no longer generates the array on the host. This can lead to less precise outputs for narrow-width floats (e.g., bfloat16) compared to previous versions.
Fix: To recover the previous host-based precision for narrow-width floats, explicitly cast to `jnp.array(np.arange(...))`.
breaking The minimum supported NumPy version for JAX is now 2.0.
Fix: Upgrade NumPy to version 2.0 or newer (and SciPy to 1.13 or newer if used).
deprecated `jax.lax.pvary` has been deprecated.
Fix: Use `jax.lax.pcast(..., to='varying')` as the replacement.

Install

pip install jax-cuda12-plugin Basic installation
pip install -U --pre jax jaxlib "jax-cuda12-plugin[with-cuda]" jax-cuda12-pjrt -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html Recommended installation with JAX core (latest pre-release)

Imports

jax
```
import jax
```
jax.numpy
```
import jax.numpy as jnp
```
While JAX's numpy API is similar, directly aliasing standard numpy as `jnp` is incorrect and will not leverage JAX's accelerator capabilities.

Quickstart

This quickstart demonstrates core JAX functionalities: utilizing the NumPy-like API (`jax.numpy`), applying Just-In-Time (JIT) compilation with `jax.jit` for performance, and computing gradients automatically using `jax.grad`. It also highlights the immutability of JAX arrays, a key difference from NumPy.

import jax
import jax.numpy as jnp

# Verify GPU device availability
print("Available devices:", jax.devices())

# Define a simple function
def f(x):
  return jnp.sum(x**2 + 2*x + 1)

# Just-in-Time compilation for performance
f_jit = jax.jit(f)

# Automatic differentiation for gradients
grad_f = jax.grad(f)
grad_f_jit = jax.jit(grad_f)

x = jnp.array([1.0, 2.0, 3.0])

print("Original function output:", f(x))
print("JIT compiled function output:", f_jit(x))
print("Gradient of function:", grad_f(x))
print("JIT compiled gradient:", grad_f_jit(x))

# Example of immutability (common gotcha):
# Attempting x[0] = 5.0 would raise a TypeError.
# Correct way to 'update' an array (creates a new array):
x_new = x.at[0].set(5.0)
print("Original array (unchanged):", x)
print("Updated array (new object):", x_new)

view raw JSON →