EnzymeAD / EnzymeAD/Enzyme.jl

Problems with callable `struct`

Open
#700 19 comments 0 reactions 0 assignees View on GitHub
activity-analysis
Dominant language
Julia
Stars
586
Forks
108
Avg merge
1d 5h
Merged PRs (30d)
44

Description

I have encountered some problems with callable `structs`, with fields being modified or wrong gradients being returned:
```julia
using AbstractGPs, TemporalGPs, Enzyme, Zygote

struct Loss
x::Vector{Float64}
y::Vector{Float64}
end

function (l::Loss)(θ)
f = to_sde(GP(θ.v * Matern52Kernel() ∘ ScaleTransform(θ.l)), SArrayStorage(Float64))
return logpdf(f(l.x, θ.σ + 1e-6), l.y)
end

θ = (v = 1., l = 1., σ = 0.1)

# First example: even though `loss` is marked `Const` (IIUC this is optional), its fields get modified:
loss = Loss(1:10, randn(10))
ref = first(loss.y)
autodiff(Reverse, Const(loss), Active(θ))
ref == first(loss.y) # false

# Second example: we introduce a shadow `struct`
# This gets rid of the problem in example 1, but the gradient is wrong
loss = Loss(1:10, randn(10))
dloss = Loss(zeros(10), zeros(10))
ref = first(loss.y)
grad_e = only(autodiff(Reverse, Duplicated(loss, dloss), Active(θ))) # needs another `|> only` for Enzyme@0.11
grad_z = only(Zygote.gradient(loss, θ))
ref == first(loss.y) # true
mapreduce(≈, &, grad_e, grad_z) # false
```

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.