EnzymeAD / Enzyme.jl

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

No method matching Duplicated(::Base.MappingRF{...}, ::Base.RefValue{Base.MappingRF{…}}) #1812

Open mhauru opened 2 months ago

mhauru commented 2 months ago

MWE:

using Random: randn
using Enzyme: Enzyme
using Turing: Turing

Enzyme.API.runtimeActivity!(true)

Turing.@model function MvDirichletWithManualAccumulation(w, doc)
    β ~ Turing.filldist(Turing.Dirichlet([1.0, 1.0]), 2)
    log_product = log.(β)
    Turing.@addlogprob! sum(log_product[CartesianIndex.(w, doc)])
end

model = MvDirichletWithManualAccumulation([1, 1, 1, 1], [1, 1, 2, 2])
ctx = Turing.DefaultContext()
vi = Turing.VarInfo(model)
vi_linked = Turing.link(vi, model)
ldp = Turing.LogDensityFunction(vi_linked, model, ctx)
f = Base.Fix1(Turing.LogDensityProblems.logdensity, ldp)
d = Turing.LogDensityProblems.dimension(ldp)
x = randn(d)

Enzyme.gradient(Enzyme.Reverse, Enzyme.Const(f), x)

Output:

ERROR: LoadError: MethodError: no method matching EnzymeCore.Duplicated(::Base.MappingRF{Base.Fix1{…}, Base.BottomRF{…}}, ::Base.RefValue{Base.MappingRF{…}})

Closest candidates are:
  EnzymeCore.Duplicated(::T1, ::T1) where T1
   @ EnzymeCore ~/.julia/packages/EnzymeCore/jSKtE/src/EnzymeCore.jl:66
  EnzymeCore.Duplicated(::T1, ::T1, ::Bool) where T1
   @ EnzymeCore ~/.julia/packages/EnzymeCore/jSKtE/src/EnzymeCore.jl:66

Stacktrace:
  [1] runtime_generic_augfwd(activity::Type{…}, width::Val{…}, ModifiedBetween::Val{…}, RT::Val{…}, f::Base.MappingRF{…}, df::Base.RefValue{…}, primal_1::Base._InitialValue, shadow_1_1::Nothing, primal_2::SubArray{…}, shadow_2_1::SubArray{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/TiboG/src/rules/jitrules.jl:338
  [2] _foldl_impl
    @ ./reduce.jl:58
  [3] foldl_impl
    @ ./reduce.jl:48 [inlined]
  [4] foldl_impl
    @ ./reduce.jl:0 [inlined]
  [5] augmented_julia_foldl_impl_8766_inner_1wrap
    @ ./reduce.jl:0
  [6] macro expansion
    @ ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:7187 [inlined]
  [7] enzyme_call
    @ ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:6794 [inlined]
  [8] AugmentedForwardThunk
    @ ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:6682 [inlined]
  [9] runtime_generic_augfwd(activity::Type{…}, width::Val{…}, ModifiedBetween::Val{…}, RT::Val{…}, f::typeof(Base.foldl_impl), df::Nothing, primal_1::Base.MappingRF{…}, shadow_1_1::Base.RefValue{…}, primal_2::Base._InitialValue, shadow_2_1::Nothing, primal_3::Distributions.EachVariate{…}, shadow_3_1::Distributions.EachVariate{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/TiboG/src/rules/jitrules.jl:338
 [10] mapfoldl_impl
    @ ./reduce.jl:44 [inlined]
 [11] mapfoldl
    @ ./reduce.jl:175 [inlined]
 [12] _mapreduce
    @ ./reduce.jl:453 [inlined]
 [13] _mapreduce_dim
    @ ./reducedim.jl:365 [inlined]
 [14] mapreduce
    @ ./reducedim.jl:357 [inlined]
 [15] #_sum#831
    @ ./reducedim.jl:1015 [inlined]
 [16] augmented_julia___sum_831_8532_inner_1wrap
    @ ./reducedim.jl:0
 [17] macro expansion
    @ ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:7187 [inlined]
 [18] enzyme_call
    @ ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:6794 [inlined]
 [19] AugmentedForwardThunk
    @ ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:6682 [inlined]
 [20] runtime_generic_augfwd(activity::Type{…}, width::Val{…}, ModifiedBetween::Val{…}, RT::Val{…}, f::Base.var"##_sum#831", df::Nothing, primal_1::@Kwargs{}, shadow_1_1::Nothing, primal_2::typeof(Base._sum), shadow_2_1::Nothing, primal_3::Base.Fix1{…}, shadow_3_1::Base.RefValue{…}, primal_4::Distributions.EachVariate{…}, shadow_4_1::Distributions.EachVariate{…}, primal_5::Colon, shadow_5_1::Nothing)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/TiboG/src/rules/jitrules.jl:338
 [21] _sum
    @ ./reducedim.jl:1015 [inlined]
 [22] sum
    @ ./reducedim.jl:1011 [inlined]
 [23] loglikelihood
    @ ~/.julia/packages/Distributions/nmUhn/src/common.jl:458 [inlined]
 [24] _logpdf
    @ ~/.julia/packages/DistributionsAD/1Mxge/src/filldist.jl:97 [inlined]
 [25] logpdf
    @ ~/.julia/packages/Distributions/nmUhn/src/common.jl:263 [inlined]
 [26] invlink_with_logpdf
    @ ~/.julia/packages/DynamicPPL/DvdZw/src/abstract_varinfo.jl:856
 [27] invlink_with_logpdf
    @ ~/.julia/packages/DynamicPPL/DvdZw/src/abstract_varinfo.jl:850 [inlined]
 [28] assume
    @ ~/.julia/packages/DynamicPPL/DvdZw/src/context_implementations.jl:227 [inlined]
 [29] tilde_assume
    @ ~/.julia/packages/DynamicPPL/DvdZw/src/context_implementations.jl:60 [inlined]
 [30] tilde_assume
    @ ~/.julia/packages/DynamicPPL/DvdZw/src/context_implementations.jl:57 [inlined]
 [31] tilde_assume!!
    @ ~/.julia/packages/DynamicPPL/DvdZw/src/context_implementations.jl:144
 [32] MvDirichletWithManualAccumulation
    @ ~/projects/Enzyme-mwes/MappingRF_BottomRF/mwe.jl:11
 [33] _evaluate!!
    @ ~/.julia/packages/DynamicPPL/DvdZw/src/model.jl:973 [inlined]
 [34] evaluate_threadunsafe!!
    @ ~/.julia/packages/DynamicPPL/DvdZw/src/model.jl:946 [inlined]
 [35] evaluate!!
    @ ~/.julia/packages/DynamicPPL/DvdZw/src/model.jl:894 [inlined]
 [36] logdensity
    @ ~/.julia/packages/DynamicPPL/DvdZw/src/logdensityfunction.jl:140
 [37] Fix1
    @ ./operators.jl:1118 [inlined]
 [38] Fix1
    @ ./operators.jl:0 [inlined]
 [39] diffejulia_Fix1_10564_inner_2wrap
    @ ./operators.jl:0
 [40] macro expansion
    @ ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:7187 [inlined]
 [41] enzyme_call
    @ ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:6794 [inlined]
 [42] CombinedAdjointThunk
    @ ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:6671 [inlined]
 [43] autodiff
    @ ~/.julia/packages/Enzyme/TiboG/src/Enzyme.jl:320 [inlined]
 [44] gradient(rm::EnzymeCore.ReverseMode{…}, f::EnzymeCore.Const{…}, x::Vector{…})
    @ Enzyme ~/.julia/packages/Enzyme/TiboG/src/Enzyme.jl:1049
 [45] top-level scope
    @ ~/projects/Enzyme-mwes/MappingRF_BottomRF/mwe.jl:25
 [46] include(fname::String)
    @ Base.MainInclude ./client.jl:489
 [47] top-level scope
    @ REPL[3]:1
in expression starting at /Users/mhauru/projects/Enzyme-mwes/MappingRF_BottomRF/mwe.jl:1
Some type information was truncated. Use `show(err)` to see complete types.
wsmoses commented 2 months ago

This one is a bit too complicated to debug directly, can you create a MWE that removes as much of tiring as possible (and ideally contains direct function defns without macros)

On Thu, Sep 12, 2024, 7:09 AM Markus Hauru @.***> wrote:

MWE:

using Random: randnusing Enzyme: Enzymeusing Turing: Turing

Enzyme.API.runtimeActivity!(true)

@. function MvDirichletWithManualAccumulation(w, doc) β ~ Turing.filldist(Turing.Dirichlet([1.0, 1.0]), 2) log_product = log.(β) @.! sum(log_product[CartesianIndex.(w, doc)])end

model = MvDirichletWithManualAccumulation([1, 1, 1, 1], [1, 1, 2, 2]) ctx = Turing.DefaultContext() vi = Turing.VarInfo(model) vi_linked = Turing.link(vi, model) ldp = Turing.LogDensityFunction(vi_linked, model, ctx) f = Base.Fix1(Turing.LogDensityProblems.logdensity, ldp) d = Turing.LogDensityProblems.dimension(ldp) x = randn(d)

Enzyme.gradient(Enzyme.Reverse, Enzyme.Const(f), x)

Output:

ERROR: LoadError: MethodError: no method matching EnzymeCore.Duplicated(::Base.MappingRF{Base.Fix1{…}, Base.BottomRF{…}}, ::Base.RefValue{Base.MappingRF{…}})

Closest candidates are: EnzymeCore.Duplicated(::T1, ::T1) where T1 @ EnzymeCore ~/.julia/packages/EnzymeCore/jSKtE/src/EnzymeCore.jl:66 EnzymeCore.Duplicated(::T1, ::T1, ::Bool) where T1 @ EnzymeCore ~/.julia/packages/EnzymeCore/jSKtE/src/EnzymeCore.jl:66

Stacktrace: [1] runtime_generic_augfwd(activity::Type{…}, width::Val{…}, ModifiedBetween::Val{…}, RT::Val{…}, f::Base.MappingRF{…}, df::Base.RefValue{…}, primal_1::Base._InitialValue, shadow_1_1::Nothing, primal_2::SubArray{…}, shadow_2_1::SubArray{…}) @ Enzyme.Compiler ~/.julia/packages/Enzyme/TiboG/src/rules/jitrules.jl:338 [2] _foldl_impl @ ./reduce.jl:58 [3] foldl_impl @ ./reduce.jl:48 [inlined] [4] foldl_impl @ ./reduce.jl:0 [inlined] [5] augmented_julia_foldl_impl_8766_inner_1wrap @ ./reduce.jl:0 [6] macro expansion @ ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:7187 [inlined] [7] enzyme_call @ ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:6794 [inlined] [8] AugmentedForwardThunk @ ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:6682 [inlined] [9] runtime_generic_augfwd(activity::Type{…}, width::Val{…}, ModifiedBetween::Val{…}, RT::Val{…}, f::typeof(Base.foldl_impl), df::Nothing, primal_1::Base.MappingRF{…}, shadow_1_1::Base.RefValue{…}, primal_2::Base._InitialValue, shadow_2_1::Nothing, primal_3::Distributions.EachVariate{…}, shadow_3_1::Distributions.EachVariate{…}) @ Enzyme.Compiler ~/.julia/packages/Enzyme/TiboG/src/rules/jitrules.jl:338 [10] mapfoldl_impl @ ./reduce.jl:44 [inlined] [11] mapfoldl @ ./reduce.jl:175 [inlined] [12] _mapreduce @ ./reduce.jl:453 [inlined] [13] _mapreduce_dim @ ./reducedim.jl:365 [inlined] [14] mapreduce @ ./reducedim.jl:357 [inlined] [15] #_sum#831 @ ./reducedim.jl:1015 [inlined] [16] augmented_julia___sum_831_8532_inner_1wrap @ ./reducedim.jl:0 [17] macro expansion @ ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:7187 [inlined] [18] enzyme_call @ ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:6794 [inlined] [19] AugmentedForwardThunk @ ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:6682 [inlined] [20] runtime_generic_augfwd(activity::Type{…}, width::Val{…}, ModifiedBetween::Val{…}, RT::Val{…}, f::Base.var"##_sum#831", df::Nothing, @.***{}, shadow_1_1::Nothing, primal_2::typeof(Base._sum), shadow_2_1::Nothing, primal_3::Base.Fix1{…}, shadow_3_1::Base.RefValue{…}, primal_4::Distributions.EachVariate{…}, shadow_4_1::Distributions.EachVariate{…}, primal_5::Colon, shadow_5_1::Nothing) @ Enzyme.Compiler ~/.julia/packages/Enzyme/TiboG/src/rules/jitrules.jl:338 [21] _sum @ ./reducedim.jl:1015 [inlined] [22] sum @ ./reducedim.jl:1011 [inlined] [23] loglikelihood @ ~/.julia/packages/Distributions/nmUhn/src/common.jl:458 [inlined] [24] _logpdf @ ~/.julia/packages/DistributionsAD/1Mxge/src/filldist.jl:97 [inlined] [25] logpdf @ ~/.julia/packages/Distributions/nmUhn/src/common.jl:263 [inlined] [26] invlink_with_logpdf @ ~/.julia/packages/DynamicPPL/DvdZw/src/abstract_varinfo.jl:856 [27] invlink_with_logpdf @ ~/.julia/packages/DynamicPPL/DvdZw/src/abstract_varinfo.jl:850 [inlined] [28] assume @ ~/.julia/packages/DynamicPPL/DvdZw/src/context_implementations.jl:227 [inlined] [29] tilde_assume @ ~/.julia/packages/DynamicPPL/DvdZw/src/context_implementations.jl:60 [inlined] [30] tilde_assume @ ~/.julia/packages/DynamicPPL/DvdZw/src/context_implementations.jl:57 [inlined] [31] tilde_assume!! @ ~/.julia/packages/DynamicPPL/DvdZw/src/context_implementations.jl:144 [32] MvDirichletWithManualAccumulation @ ~/projects/Enzyme-mwes/MappingRF_BottomRF/mwe.jl:11 [33] _evaluate!! @ ~/.julia/packages/DynamicPPL/DvdZw/src/model.jl:973 [inlined] [34] evaluate_threadunsafe!! @ ~/.julia/packages/DynamicPPL/DvdZw/src/model.jl:946 [inlined] [35] evaluate!! @ ~/.julia/packages/DynamicPPL/DvdZw/src/model.jl:894 [inlined] [36] logdensity @ ~/.julia/packages/DynamicPPL/DvdZw/src/logdensityfunction.jl:140 [37] Fix1 @ ./operators.jl:1118 [inlined] [38] Fix1 @ ./operators.jl:0 [inlined] [39] diffejulia_Fix1_10564_inner_2wrap @ ./operators.jl:0 [40] macro expansion @ ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:7187 [inlined] [41] enzyme_call @ ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:6794 [inlined] [42] CombinedAdjointThunk @ ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:6671 [inlined] [43] autodiff @ ~/.julia/packages/Enzyme/TiboG/src/Enzyme.jl:320 [inlined] [44] gradient(rm::EnzymeCore.ReverseMode{…}, f::EnzymeCore.Const{…}, x::Vector{…}) @ Enzyme ~/.julia/packages/Enzyme/TiboG/src/Enzyme.jl:1049 [45] top-level scope @ ~/projects/Enzyme-mwes/MappingRF_BottomRF/mwe.jl:25 [46] include(fname::String) @ Base.MainInclude ./client.jl:489 [47] top-level scope @ REPL[3]:1 in expression starting at /Users/mhauru/projects/Enzyme-mwes/MappingRF_BottomRF/mwe.jl:1 Some type information was truncated. Use show(err) to see complete types.

— Reply to this email directly, view it on GitHub https://github.com/EnzymeAD/Enzyme.jl/issues/1812, or unsubscribe https://github.com/notifications/unsubscribe-auth/AAJTUXE3IJ7GVLTUVGLTTP3ZWFY7RAVCNFSM6AAAAABOC7X44OVHI2DSMVQWIX3LMV43ASLTON2WKOZSGUZDEMBYHE3TMNY . You are receiving this because you are subscribed to this thread.Message ID: @.***>

mhauru commented 2 months ago

Somewhat simplified:

using Random: randn
using Enzyme: Enzyme
using DynamicPPL: DynamicPPL
using Distributions: Dirichlet
using DistributionsAD: filldist

Enzyme.API.runtimeActivity!(true)

w = [1, 1, 1, 1]
doc = [1, 1, 2, 2]
vn = DynamicPPL.VarName{:β}()
vi_template = DynamicPPL.VarInfo()
d = filldist(Dirichlet([1.0, 1.0]), 2)
β = rand(2,2)
vi = DynamicPPL.push!!(vi_template, vn, β, d, Set{DynamicPPL.Selector}())

function f(x)
    vi = DynamicPPL.unflatten(vi_template, x)
    ctx = DynamicPPL.DefaultContext()
    β, vi = DynamicPPL.tilde_assume!!(
        ctx, filldist(Dirichlet([1.0, 1.0]), 2), vn, vi
    )
    log_product = log.(β)
    vi = DynamicPPL.acclogp!!(ctx, vi, sum(log_product[CartesianIndex.(w, doc)]))
    return DynamicPPL.getlogp(vi)
end

x = randn(4)

Enzyme.gradient(Enzyme.Reverse, Enzyme.Const(f), x)
wsmoses commented 2 months ago

Hm sorry, but could you give it a bit more minimization still?

wsmoses commented 2 months ago

actually wait, nevermind

mhauru commented 2 months ago

Can do more if it helps, the above is just what I got to yesterday before calling it a day.

wsmoses commented 1 month ago

With the DynamicPPL release out, closing this

mhauru commented 1 month ago

Sorry @wsmoses, I don't understand. This still fails with the same error on the latest DynamicPPL and Enzyme main. Reopen?

Above MWE with the new way of setting runtime activity:

using Random: randn
using Enzyme: Enzyme
using DynamicPPL: DynamicPPL
using Distributions: Dirichlet
using DistributionsAD: filldist

w = [1, 1, 1, 1]
doc = [1, 1, 2, 2]
vn = DynamicPPL.VarName{:β}()
vi_template = DynamicPPL.VarInfo()
d = filldist(Dirichlet([1.0, 1.0]), 2)
β = rand(2,2)
vi = DynamicPPL.push!!(vi_template, vn, β, d, Set{DynamicPPL.Selector}())

function f(x)
    vi = DynamicPPL.unflatten(vi_template, x)
    ctx = DynamicPPL.DefaultContext()
    β, vi = DynamicPPL.tilde_assume!!(
        ctx, filldist(Dirichlet([1.0, 1.0]), 2), vn, vi
    )
    log_product = log.(β)
    vi = DynamicPPL.acclogp!!(ctx, vi, sum(log_product[CartesianIndex.(w, doc)]))
    return DynamicPPL.getlogp(vi)
end

x = randn(4)

Enzyme.gradient(Enzyme.set_runtime_activity(Enzyme.Reverse), Enzyme.Const(f), x)
wsmoses commented 1 month ago

oh sorry misremembered this one and the other one

wsmoses commented 1 month ago

and yeah okay sadly this needs further reduction as well

mhauru commented 1 month ago

Reduced:

module MWE

using Enzyme: Enzyme
using Distributions: Distributions

Enzyme.API.runtimeActivity!(true)

function f(x)
    d = Distributions.Dirichlet([1.0, 1.0])
    v = eachrow(x)
    return sum(Base.Fix1(Distributions.logpdf, d), v)
end

x = reshape([0.8, 0.3, 0.2, 0.7], (2, 2))
@show f(x)
Enzyme.gradient(Enzyme.Reverse, f, x)

end
wsmoses commented 1 month ago

I think we’re going to need it even further reduced here.

Also I don’t think you’re on the latest Enzyme since you’re using an APi call which was removed

On Fri, Oct 4, 2024 at 5:01 AM Markus Hauru @.***> wrote:

Reduced:

module MWE using Enzyme: Enzymeusing Distributions: Distributions

Enzyme.API.runtimeActivity!(true) function f(x) d = Distributions.Dirichlet([1.0, 1.0]) v = eachrow(x) return sum(Base.Fix1(Distributions.logpdf, d), v)end

x = reshape([0.8, 0.3, 0.2, 0.7], (2, @.*** f(x) Enzyme.gradient(Enzyme.Reverse, f, x) end

— Reply to this email directly, view it on GitHub https://github.com/EnzymeAD/Enzyme.jl/issues/1812#issuecomment-2393331059, or unsubscribe https://github.com/notifications/unsubscribe-auth/AAJTUXHGAACCKH64XUDDXMLZZZRQHAVCNFSM6AAAAABOC7X44OVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDGOJTGMZTCMBVHE . You are receiving this because you modified the open/close state.Message ID: @.***>

mhauru commented 1 month ago

Same error on Enzyme v0.13.7 (once you remove the runtimeActivity! setting).

I gotta go attend to other things, feel free to minimise further if you have time.

wsmoses commented 2 weeks ago

@penelopeysm if you have time, would you be able to reduce this one as well?