Change to DifferentiationInterface #46
Conversation
|
Looks good to me! |
| ), | ||
| ) | ||
| time_elapsed = @elapsed begin | ||
| opt_stats = [] |
There was a problem hiding this comment.
If you want to avoid the abstract container (and map(identity,...) below) you could move the first iteration out of the loop. Or define it unitialized with local opt_stats here and set it in the i = 1 iteration, similar to
julia> function f(n::Int)
local out
i = 1
while i <= n
out = i == 1 ? [1] : [i + only(out)]
i += 1
end
return only(out)
end
julia> f(3)
6
| ADTypes.AutoReverseDiff(false), | ||
| ADTypes.AutoEnzyme(), | ||
| ADTypes.AutoReverseDiff(; compile=false), | ||
| ADTypes.AutoMooncake(; config=Mooncake.Config()), |
There was a problem hiding this comment.
Why remove the test with Enzyme?
| ADTypes.AutoReverseDiff(false), | ||
| # ADTypes.AutoEnzyme(), # not working now | ||
| ADTypes.AutoReverseDiff(; compile = false), | ||
| ADTypes.AutoMooncake(; config=Mooncake.Config()), |
There was a problem hiding this comment.
The test was working with Enzyme (and runtime activity) in https://github.com/TuringLang/NormalizingFlows.jl/pull/38/files.
There was a problem hiding this comment.
Maybe try with AutoEnzyme(mode=Enzyme.set_runtime_activity(Enzyme.Reverse))?
|
@zuhengxu PR looks great, I can take over and finish integrating the reviews if you are busy |
|
Thank you all for the review and comments!
Hi @sunxd3, I'm bit busy before Thursday this week. It will be much appreciated if you can help to integrate some of the review! Otherwise I'll integrate them by this Friday. |
Co-authored-by: David Widmann <[email protected]>
Change the optimization interface to DifferentatiationInterface for easier switch/test between different AD systems.