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
In [1]:
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