jax-cuda12-plugin

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)