EnzymeAD / Enzyme.jl

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

BatchDuplicated fails with RefValues #1997

Open mhauru opened 1 week ago

mhauru commented 1 week ago

Feel free to rename the issue to be more informative, I don't really understand what's going on here.

MWE:

julia> import Enzyme; using Distributions

julia> func = x -> logpdf(reshape(product_distribution([Normal(), Uniform()]), 1, 2), x)

#27 (generic function with 1 method)

julia> Enzyme.gradient(Enzyme.Forward, Enzyme.Const(func), [2.1 0.7])
ERROR: MethodError: no method matching EnzymeCore.BatchDuplicated(::Normal{Float64}, ::Tuple{Base.RefValue{Normal{Float64}}, Base.RefValue{Normal{Float64}}})

Closest candidates are:
  EnzymeCore.BatchDuplicated(::T1, ::Tuple{Vararg{T1, N}}) where {T1, N}
   @ EnzymeCore ~/projects/Enzyme.jl/lib/EnzymeCore/src/EnzymeCore.jl:134
  EnzymeCore.BatchDuplicated(::T1, ::Tuple{Vararg{T1, N}}, ::Bool) where {T1, N}
   @ EnzymeCore ~/projects/Enzyme.jl/lib/EnzymeCore/src/EnzymeCore.jl:134

Stacktrace:
  [1] runtime_generic_fwd(activity::Type{…}, runtimeActivity::Val{…}, width::Val{…}, RT::Val{…}, f::typeof(logpdf), df::Nothing, df_2::Nothing, primal_1::Normal{…}, shadow_1_1::Base.RefValue{…}, shadow_1_2::Base.RefValue{…}, primal_2::Float64, shadow_2_1::Float64, shadow_2_2::Float64)
    @ Enzyme.Compiler ~/projects/Enzyme.jl/src/rules/jitrules.jl:170
  [2] #533
    @ ~/.julia/packages/Distributions/nmUhn/src/multivariate/product.jl:46
  [3] _mapreduce
    @ ./reduce.jl:440
  [4] _mapreduce_dim
    @ ./reducedim.jl:365 [inlined]
  [5] mapreduce
    @ ./reducedim.jl:357 [inlined]
  [6] _sum
    @ ./reducedim.jl:1015 [inlined]
  [7] sum
    @ ./reducedim.jl:1011 [inlined]
  [8] _logpdf
    @ ~/.julia/packages/Distributions/nmUhn/src/multivariate/product.jl:46
  [9] logpdf
    @ ~/.julia/packages/Distributions/nmUhn/src/common.jl:263 [inlined]
 [10] __logpdf
    @ ~/.julia/packages/Distributions/nmUhn/src/reshaped.jl:55 [inlined]
 [11] _logpdf
    @ ~/.julia/packages/Distributions/nmUhn/src/reshaped.jl:52 [inlined]
 [12] logpdf
    @ ~/.julia/packages/Distributions/nmUhn/src/common.jl:263 [inlined]
 [13] #27
    @ ./REPL[126]:1 [inlined]
 [14] fwddiffe2julia__27_208218wrap
    @ ./REPL[126]:0
 [15] macro expansion
    @ ~/projects/Enzyme.jl/src/compiler.jl:8148 [inlined]
 [16] enzyme_call
    @ ~/projects/Enzyme.jl/src/compiler.jl:7714 [inlined]
 [17] ForwardModeThunk
    @ ~/projects/Enzyme.jl/src/compiler.jl:7503 [inlined]
 [18] autodiff
    @ ~/projects/Enzyme.jl/src/Enzyme.jl:647 [inlined]
 [19] macro expansion
    @ ~/projects/Enzyme.jl/src/Enzyme.jl:2068 [inlined]
 [20] gradient(::EnzymeCore.ForwardMode{…}, ::EnzymeCore.Const{…}, ::Matrix{…}; chunk::Nothing, shadows::Tuple{…})
    @ Enzyme ~/projects/Enzyme.jl/src/Enzyme.jl:1970
 [21] gradient(::EnzymeCore.ForwardMode{…}, ::EnzymeCore.Const{…}, ::Matrix{…})
    @ Enzyme ~/projects/Enzyme.jl/src/Enzyme.jl:1970
 [22] top-level scope
    @ REPL[128]:1
Some type information was truncated. Use `show(err)` to see complete types.
wsmoses commented 1 week ago

This is quite confusing, a further reduction would be helpful here

mhauru commented 1 week ago

Would like to help, but unfortunately I'm too busy right now to find time to minimise further. Open for grabs for anyone to do it.