search

JuliaDiff / ReverseDiff.jl

Reverse Mode Automatic Differentiation for Julia
Other
348 stars 57 forks source link

Strange bug when deferring to ChainRules #236

Open torfjelde opened 1 year ago

torfjelde commented 1 year ago

I unfortunately haven't been able to debug this properly yet, but I encountered a somewhat strange bug when using @grad_from_chainrules.

In https://github.com/TuringLang/Bijectors.jl/actions/runs/5781513384/job/15666726634#step:5:274 you can see

   BoundsError: attempt to access 4×4 Matrix{Float64} at index [0]
  Stacktrace:
    [1] getindex
      @ ./array.jl:805 [inlined]
    [2] getindex
      @ ./multidimensional.jl:643 [inlined]
    [3] special_reverse_exec!(instruction::ReverseDiff.SpecialInstruction{Tuple{typeof(getindex), Val{:generic}}, Tuple{ReverseDiff.TrackedArray{Float64, Float64, 2, Matrix{Float64}, Matrix{Float64}}, Tuple{Matrix{Bool}}}, ReverseDiff.TrackedArray{Float64, Float64, 1, Vector{Float64}, Vector{Float64}}, Nothing})
      @ ReverseDiff ~/.julia/packages/ReverseDiff/7pHoq/src/tracked.jl:372
    [4] reverse_exec!(instruction::ReverseDiff.SpecialInstruction{Tuple{typeof(getindex), Val{:generic}}, Tuple{ReverseDiff.TrackedArray{Float64, Float64, 2, Matrix{Float64}, Matrix{Float64}}, Tuple{Matrix{Bool}}}, ReverseDiff.TrackedArray{Float64, Float64, 1, Vector{Float64}, Vector{Float64}}, Nothing})
      @ ReverseDiff ~/.julia/packages/ReverseDiff/7pHoq/src/tape.jl:93
    [5] reverse_pass!(tape::Vector{ReverseDiff.AbstractInstruction})
      @ ReverseDiff ~/.julia/packages/ReverseDiff/7pHoq/src/tape.jl:87
    [6] reverse_pass!
      @ ~/.julia/packages/ReverseDiff/7pHoq/src/api/tape.jl:36 [inlined]
    [7] seeded_reverse_pass!(result::Vector{Float64}, output::ReverseDiff.TrackedReal{Float64, Float64, Nothing}, input::ReverseDiff.TrackedArray{Float64, Float64, 1, Vector{Float64}, Vector{Float64}}, tape::ReverseDiff.GradientTape{var"#31#33"{Bijectors.PDVecBijector, Int64}, ReverseDiff.TrackedArray{Float64, Float64, 1, Vector{Float64}, Vector{Float64}}, ReverseDiff.TrackedReal{Float64, Float64, Nothing}})
      @ ReverseDiff ~/.julia/packages/ReverseDiff/7pHoq/src/api/utils.jl:31
    [8] seeded_reverse_pass!(result::Vector{Float64}, t::ReverseDiff.GradientTape{var"#31#33"{Bijectors.PDVecBijector, Int64}, ReverseDiff.TrackedArray{Float64, Float64, 1, Vector{Float64}, Vector{Float64}}, ReverseDiff.TrackedReal{Float64, Float64, Nothing}})
      @ ReverseDiff ~/.julia/packages/ReverseDiff/7pHoq/src/api/tape.jl:47
    [9] gradient(f::Function, input::Vector{Float64}, cfg::ReverseDiff.GradientConfig{ReverseDiff.TrackedArray{Float64, Float64, 1, Vector{Float64}, Vector{Float64}}})
      @ ReverseDiff ~/.julia/packages/ReverseDiff/7pHoq/src/api/gradients.jl:24
   [10] gradient
      @ ~/.julia/packages/ReverseDiff/7pHoq/src/api/gradients.jl:22 [inlined]
   [11] test_ad(f::Function, x::Vector{Float64}, broken::Tuple{}; rtol::Float64, atol::Float64, use_forwarddiff_as_truth::Bool)
      @ Main ~/work/Bijectors.jl/Bijectors.jl/test/ad/utils.jl:33
   [12] macro expansion
      @ ~/work/Bijectors.jl/Bijectors.jl/test/ad/pd.jl:9 [inlined]
   [13] macro expansion
      @ /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.6/Test/src/Test.jl:1151 [inlined]
   [14] top-level scope
      @ ~/work/Bijectors.jl/Bijectors.jl/test/ad/pd.jl:2

If I then replace the offending rule with a "hard-coded" version using ChainRules, i.e. apply https://github.com/TuringLang/Bijectors.jl/pull/280/commits/29790dcd12e659d2f311b92b8f75a72e03802f5e, then the above error goes away :confused:

I don't have time to dig into this right now, but figured I'd put it here for future reference.

torfjelde commented 1 year ago

Okay, so it turns out the error only shows up when using @grad_from_chainrules because we somehow hit getindex with the index being a boolean array (whose elements are then subsequently treated as integer indices in the pullback), while this does not happen when we use @grad (for some reason).

~Is this maybe because @grad_from_chainrules does not call value one the input?~

EDIT: Nope; this is done: https://github.com/JuliaDiff/ReverseDiff.jl/blob/105290ffff86a5eccc55ca64feee0a1c52fffad0/src/macros.jl#L338