EnzymeAD / Enzyme.jl

Julia bindings for the Enzyme automatic differentiator
https://enzyme.mit.edu
MIT License
422 stars 57 forks source link

Enzyme produces inconsistent results with macros from UnPack.jl #1579

Closed jlk9 closed 3 days ago

jlk9 commented 3 days ago

UnPack.jl provides functionality to get fields out of structs (@unpack) or replace fields back into structs (@pack!) with clean syntax. But it produces incorrect results with Enzyme. Running this code

using Enzyme
using UnPack

mutable struct PrognosticVars{F<:AbstractFloat, FV <: AbstractArray{F,1}}

    # var: sea surface height [m] 
    # dim: (nCells)
    a::FV

    function PrognosticVars(a::AT2D) where {AT2D}

        type = eltype(a)
        new{type, AT2D}(a)
    end        
end

function get_sum(Prog::PrognosticVars)

    sum = 0.0
    for j = 1:size(Prog.a)[1]
        sum = sum + Prog.a[j]
    end
    return sum
end

function get_sum_pack(Prog::PrognosticVars)

    @unpack a = Prog

    sum = 0.0
    for j = 1:size(a)[1]
        sum = sum + a[j]
    end

    @pack! Prog = a

    return sum
end

a = [1., 2., 3.]

prog = PrognosticVars(a)

#sum = get_sum(prog)

d_prog = Enzyme.make_zero(prog)
d_sum  = autodiff(Enzyme.Reverse, get_sum, Duplicated(prog, d_prog))
@show d_sum, d_prog

d_prog_pack = Enzyme.make_zero(prog)
d_sum_pack  = autodiff(Enzyme.Reverse, get_sum_pack, Duplicated(prog, d_prog))

@show d_sum_pack, d_prog_pack

produces these results:

(d_sum, d_prog) = (((nothing,),), PrognosticVars{Float64, Vector{Float64}}([1.0, 1.0, 1.0]))
(d_sum_pack, d_prog_pack) = (((nothing,),), PrognosticVars{Float64, Vector{Float64}}([0.0, 0.0, 0.0]))

Logically d_prog_pack should equal d_prog, since they both sum up the contents of a. The pack! and unpack macros are just for syntax convenience.

wsmoses commented 3 days ago

@jlk9 this is a bug in your test case.

In the second call you're pasing d_prog not d_prog_pack, using the correct variable works.