Failure to compute gradient for vcat(...)

[nsiccha]

There's some funky failure to compute the gradient for the following MWE using vcat. I'm assuming similar things happen for hcat/cat.

MWE

using DifferentiationInterface
import Mooncake

f(x) = vcat(
    x[1],
    g(x[2:2]),
    g(x[3:3]),
    x[4:4]
)
g(x) = x .- log(sum(exp.(x)))
h(x) = sum(f(x))
backend = AutoMooncake(; config=nothing)
x = ones(4)
prep = prepare_gradient(h, backend, x)
gradient(h, prep, backend, x)

gives me

ERROR: BoundsError: attempt to access Tuple{Mooncake.NoRData, Mooncake.NoRData, Mooncake.NoRData, Mooncake.NoRData, Mooncake.NoRData, Float64, Mooncake.NoRData, Mooncake.NoRData} at index [9]
Stacktrace:
  [1] f
    @ ~/github/nsiccha/mwe/Mooncake/index.jl:4 [inlined]
  [2] (::Tuple{…})(none::Mooncake.NoRData)
    @ Base.Experimental ./<missing>:0
  [3] Pullback
    @ ~/.julia/packages/Mooncake/nSw4x/src/interpreter/s2s_reverse_mode_ad.jl:888 [inlined]
  [4] h
    @ ~/github/nsiccha/mwe/Mooncake/index.jl:11 [inlined]
  [5] (::Tuple{…})(none::Float64)
    @ Base.Experimental ./<missing>:0
  [6] Pullback
    @ ~/.julia/packages/Mooncake/nSw4x/src/interpreter/s2s_reverse_mode_ad.jl:888 [inlined]
  [7] __value_and_pullback!!(::Mooncake.DerivedRule{…}, ::Float64, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}; y_cache::Float64)
    @ Mooncake ~/.julia/packages/Mooncake/nSw4x/src/interface.jl:20
  [8] __value_and_pullback!!
    @ ~/.julia/packages/Mooncake/nSw4x/src/interface.jl:11 [inlined]
  [9] value_and_pullback!!
    @ ~/.julia/packages/Mooncake/nSw4x/src/interface.jl:219 [inlined]
 [10] value_and_pullback(::typeof(h), ::DifferentiationInterfaceMooncakeExt.MooncakeOneArgPullbackPrep{…}, ::AutoMooncake{…}, ::Vector{…}, ::Tuple{…})
    @ DifferentiationInterfaceMooncakeExt ~/.julia/packages/DifferentiationInterface/Tls8M/ext/DifferentiationInterfaceMooncakeExt/onearg.jl:30
 [11] prepare_pullback(::Function, ::AutoMooncake{Nothing}, ::Vector{Float64}, ::Tuple{Bool})
    @ DifferentiationInterfaceMooncakeExt ~/.julia/packages/DifferentiationInterface/Tls8M/ext/DifferentiationInterfaceMooncakeExt/onearg.jl:16
 [12] prepare_gradient(::typeof(h), ::AutoMooncake{Nothing}, ::Vector{Float64})
    @ DifferentiationInterface ~/.julia/packages/DifferentiationInterface/Tls8M/src/first_order/gradient.jl:70
 [13] top-level scope
    @ ~/github/nsiccha/mwe/Mooncake/index.jl:14
Some type information was truncated. Use `show(err)` to see complete types.

Varying the arguments given to vcat sometimes removes the error. E.g. the following all works:

• f(x) = vcat(x[1], g(x[3:3]), x[4:4])
• f(x) = vcat(x[1], g(x[2:2]), g(x[3:3]), x[4])
• f(x) = vcat(x[1:1], g(x[2:2]), g(x[3:3]), x[4:4])

Version info:

julia version 1.10.8
[a0c0ee7d] DifferentiationInterface v0.6.34
[da2b9cff] Mooncake v0.4.78

[willtebbutt]

Oooooo this is very strange. Thanks for opening the issue -- I'll debug and get back to you!

[willtebbutt]

This should be resolved once version 0.4.79 is available (which should happen in the next 30 mins or so). Please re-open if it's not.