Back to Article
Smoke: Reactant
Download Source

Smoke: Reactant

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