Toolchain check

A one-line gradient with each tool, to verify your setup

using CUDA, KernelAbstractions, Enzyme
@kernel function square!(y, @Const(x))
    i = @index(Global)
    @inbounds y[i] = x[i]^2
end
f!(y, x) = (square!(get_backend(x))(y, x; ndrange=length(x)); nothing)
x = CuArray(Float32[1, 2, 3, 4]); dx = CUDA.zeros(Float32, 4)
y = CUDA.zeros(Float32, 4);       dy = CUDA.ones(Float32, 4)
Enzyme.autodiff(Reverse, f!, Const, Duplicated(y, dy), Duplicated(x, dx))
Array(dx)

4-element Vector{Float32}:
 2.0
 4.0
 6.0
 8.0

Source: Smoke: Enzyme

using Reactant, Enzyme
Reactant.set_default_backend("gpu")
loss(x) = sum(sort(x) .^ 2)
g(x) = Enzyme.gradient(Reverse, loss, x)[1]
x = Reactant.to_rarray(Float32[3, 1, 4, 2])
Array((@compile g(x))(x))

4-element Vector{Float32}:
 6.0
 2.0
 8.0
 4.0

Source: Smoke: Reactant

import os
os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false")
import jax, jax.numpy as jnp
print(jax.devices())
jax.grad(lambda x: jnp.sum(jnp.sort(x) ** 2))(jnp.array([3., 1., 4., 2.]))

[CudaDevice(id=0)]

Array([6., 2., 8., 4.], dtype=float32)

Source: Smoke: JAX

Notebooks