ForwardDiff tagging errors with MarginalLogDensities

[ElOceanografo]

I just noticed this issue, which appeared for me between DI v0.6.17 and v0.6.18: I get ForwardDiff "Invalid Tag object" errors when calling a MarginalLogDensity that uses ForwardDiff for the Hessian. MWE:

using DifferentiationInterface
using MarginalLogDensities
import Zygote

N = 3
ld(x) = sum(abs2, diff(x))
iw = [1, 3]
iv = [2]
u = randn(N)
v = u[iv]
w = u[iw]

w = rand(3)
mld1 = MarginalLogDensity(ld, u, iw, (), LaplaceApprox())
mld1(w[iv], ()) # error

ERROR: Invalid Tag object:
  Expected ForwardDiff.Tag{DifferentiationInterface.FixTail{SciMLBase.OptimizationFunction{true, AutoForwardDiff{nothing, Nothing}, MarginalLogDensities.var"#f#4"{typeof(ld), Vector{Float64}, Vector{Int64}, Vector{Int64}}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Tuple{@NamedTuple{p::Tuple{}, v::Vector{Float64}}}}, ForwardDiff.Dual{ForwardDiff.Tag{typeof(DifferentiationInterface.shuffled_gradient), Float64}, Float64, 1}},
  Observed ForwardDiff.Tag{DifferentiationInterface.FixTail{MarginalLogDensities.var"#f#4"{typeof(ld), Vector{Float64}, Vector{Int64}, Vector{Int64}}, Tuple{@NamedTuple{p::Tuple{}, v::Vector{Float64}}}}, ForwardDiff.Dual{ForwardDiff.Tag{typeof(DifferentiationInterface.shuffled_gradient), Float64}, Float64, 1}}.
Stacktrace:
  [1] checktag(::Type{…}, f::DifferentiationInterface.FixTail{…}, x::Vector{…})
    @ ForwardDiff ~/.julia/packages/ForwardDiff/UBbGT/src/config.jl:34
  [2] gradient(f::DifferentiationInterface.FixTail{…}, x::Vector{…}, cfg::ForwardDiff.GradientConfig{…}, ::Val{…})
    @ ForwardDiff ~/.julia/packages/ForwardDiff/UBbGT/src/gradient.jl:18
  [3] gradient(f::SciMLBase.OptimizationFunction{…}, prep::DifferentiationInterfaceForwardDiffExt.ForwardDiffGradientPrep{…}, backend::AutoForwardDiff{…}, x::Vector{…}, contexts::Constant{…})
    @ DifferentiationInterfaceForwardDiffExt ~/.julia/packages/DifferentiationInterface/IZ4Ns/ext/DifferentiationInterfaceForwardDiffExt/onearg.jl:357
  [4] shuffled_gradient(x::Vector{…}, f::SciMLBase.OptimizationFunction{…}, prep::DifferentiationInterfaceForwardDiffExt.ForwardDiffGradientPrep{…}, backend::AutoForwardDiff{…}, rewrap::DifferentiationInterface.Rewrap{…}, unannotated_contexts::@NamedTuple{…})
    @ DifferentiationInterface ~/.julia/packages/DifferentiationInterface/IZ4Ns/src/first_order/gradient.jl:140
  [5] compute_ydual_onearg(::typeof(DifferentiationInterface.shuffled_gradient), ::DifferentiationInterfaceForwardDiffExt.ForwardDiffOneArgPushforwardPrep{…}, ::Vector{…}, ::Tuple{…}, ::Constant{…}, ::DifferentiationInterface.PrepContext{…}, ::Constant{…}, ::Constant{…}, ::Constant{…})
    @ DifferentiationInterfaceForwardDiffExt ~/.julia/packages/DifferentiationInterface/IZ4Ns/ext/DifferentiationInterfaceForwardDiffExt/onearg.jl:106
  [6] pushforward!(::typeof(DifferentiationInterface.shuffled_gradient), ::Tuple{…}, ::DifferentiationInterfaceForwardDiffExt.ForwardDiffOneArgPushforwardPrep{…}, ::AutoForwardDiff{…}, ::Vector{…}, ::Tuple{…}, ::Constant{…}, ::DifferentiationInterface.PrepContext{…}, ::Constant{…}, ::Constant{…}, ::Constant{…})
    @ DifferentiationInterfaceForwardDiffExt ~/.julia/packages/DifferentiationInterface/IZ4Ns/ext/DifferentiationInterfaceForwardDiffExt/onearg.jl:161
  [7] hvp!(f::SciMLBase.OptimizationFunction{…}, tg::Tuple{…}, prep::DifferentiationInterfaceForwardDiffExt.ForwardDiffOverSomethingHVPPrep{…}, backend::SecondOrder{…}, x::Vector{…}, tx::Tuple{…}, contexts::Constant{…})
    @ DifferentiationInterfaceForwardDiffExt ~/.julia/packages/DifferentiationInterface/IZ4Ns/ext/DifferentiationInterfaceForwardDiffExt/secondorder.jl:70
  [8] hessian!
    @ ~/.julia/packages/DifferentiationInterface/IZ4Ns/ext/DifferentiationInterfaceSparseMatrixColoringsExt/hessian.jl:98 [inlined]
  [9] modal_hessian!
    @ ~/.julia/packages/MarginalLogDensities/hMxUz/src/MarginalLogDensities.jl:288 [inlined]
 [10] _marginalize(mld::MarginalLogDensity{…}, v::Vector{…}, data::Tuple{}, method::LaplaceApprox{…}, verbose::Bool)
    @ MarginalLogDensities ~/.julia/packages/MarginalLogDensities/hMxUz/src/MarginalLogDensities.jl:297
 [11] #_#5
    @ ~/.julia/packages/MarginalLogDensities/hMxUz/src/MarginalLogDensities.jl:226 [inlined]
 [12] (::MarginalLogDensity{…})(v::Vector{…}, data::Tuple{})
    @ MarginalLogDensities ~/.julia/packages/MarginalLogDensities/hMxUz/src/MarginalLogDensities.jl:225
 [13] top-level scope
    @ Untitled-1:16
Some type information was truncated. Use `show(err)` to see complete types.

I get the same error if I use dense ForwardDiff for the Hessian:

mld2 = MarginalLogDensity(ld, u, iw, (), LaplaceApprox(), 
    hess_adtype=AutoForwardDiff())
mld2(w[iv], ()) 

Confirming it only happens when ForwardDiff is used for the Hessian:

mld3 = MarginalLogDensity(ld, u, iw, (), LaplaceApprox(adtype=AutoForwardDiff()), 
    hess_adtype=AutoZygote())
mld3(w[iv], ()) # works

mld4 = MarginalLogDensity(ld, u, iw, (), LaplaceApprox(adtype=AutoZygote()), 
    hess_adtype=AutoForwardDiff())
mld4(w[iv], ()) # error

# defining the SecondOrder manually to confirm defaults are doing what I think they are
ad = AutoForwardDiff()
hess_ad = AutoSparse(
            SecondOrder(AutoForwardDiff(), ad),
            DenseSparsityDetector(ad, atol=sqrt(eps())),
            GreedyColoringAlgorithm())
mld5 = MarginalLogDensity(ld, u, iw, (), LaplaceApprox(adtype=ad), 
    hess_adtype=hess_ad)
mld5(w[iv], ()) # error

[gdalle]

Thanks for reporting this, I'm gonna look into it

[gdalle]

I tried to reproduce the correct behavior by pinning DI to v0.6.17 (and cloning the MarginalLogDensities repo) but I got another error, any clue why?

EDIT: I assume it's because you forgot the p in the following line:

ld(x, p) = sum(abs2, diff(x))

[gdalle]

Okay I think I found the cause and it is on the MarginalLogDensities side. I opened ElOceanografo/MarginalLogDensities.jl#36 to fix it, can you confirm that it works?

[ElOceanografo]

Yes, looks like that worked. Thanks for fixing my bugs again!