Strange bug when deferring to ChainRules

[torfjelde]

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 TuringLang/Bijectors.jl@29790dc, then the above error goes away 😕

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

[torfjelde]

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:

ReverseDiff.jl/src/macros.jl

Line 338 in 105290f

output_value, back = ChainRulesCore.rrule($f, map(ReverseDiff.value, args)...; $kwargs...)