wsmoses
commented
1 year ago On latest main you should now get a nicer error, which can be resolved by adding the flags below. However it does run with a segfault for me.
@vchuravy
using Enzyme
#Enzyme.API.printall!(true)
# Enzyme.API.printtype!(true)
Enzyme.API.strictAliasing!(false)
Enzyme.API.looseTypeAnalysis!(true)
abstract type Node end
struct Leaf <: Node
params::NamedTuple
x::Vector{Float32}
end
function (n::Leaf)()::Tuple{Float32, Vector{Float32}}
y = n.params.w2 * n.x .+ n.params.b2 |> first
return y, n.x
end
struct Branch <: Node
params::NamedTuple
left::Node
right::Node
end
function (n::Branch)()::Tuple{Float32, Vector{Float32}}
y1, h1 = n.left()
y2, h2 = n.right()
h = n.params.w1 * vcat(h1, h2) .+ n.params.b1
y3 = n.params.w2 * h .+ n.params.b2 |> first
y = y1+y2+y3
return y, h
end
function genTree(N::Int, K::Int, ps::Vector, p::Float64)::Node
if rand() < p
return Leaf(ps[rand(1:N)], rand(Float32, K))
else
return Branch(ps[rand(1:N)], genTree(N, K, ps, p), genTree(N, K, ps, p))
end
end
function loss(t1::Node, t2::Node)
return t1()[1] - t2()[1]
end
function mirror(n::Leaf, mm::Dict)
return Leaf(mm[n.params], zero(n.x))
end
function mirror(n::Branch, mm::Dict)
return Branch(mm[n.params], mirror(n.left, mm), mirror(n.right, mm))
end
function main()
N = 10
K = 5
ps = map(1:N) do x
(
w1 = randn(Float32, (K,2*K)),
b1 = randn(Float32, K),
w2 = randn(Float32, (1,K)),
b2 = randn(Float32),
)
end
grads = map(1:N) do x
(
w1 = zeros(Float32, (K,2*K)),
b1 = zeros(Float32, K),
w2 = zeros(Float32, (1,K)),
b2 = zeros(Float32),
)
end
mm = Dict(p=>g for (p, g) in zip(ps, grads))
# create a tree
t1 = genTree(N, K, ps, 0.5)
t2 = genTree(N, K, ps, 0.5)
t1Grads = mirror(t1, mm)
t2Grads = mirror(t2, mm)
@show ll = loss(t1, t2)
@show loss(t1Grads, t2Grads)
autodiff(Reverse, loss, Duplicated(t1, t1Grads), Duplicated(t2, t2Grads))
@show ll = loss(t1, t2)
end
main()
freddycct
An example of a Recursive Neural Network on Enzyme (state-of-the-art LLVM Auto Diff).
error