Troubles with ForwardColorJacCache

[briochemc]

It seems using ForwardColorJacCache does not work with anonymous functions or when the function fed to it is different from the function fed to forwarddiff_color_jacobian!. (Sorry I couldn't think of anything more specific for the title.)

Minimal example

Below, I'm simply trying to get the Jacobian of a (in-place) linear function x -> jac * x:

using SparseDiffTools, LinearAlgebra # SparseDiffTools v1.13.0
nt, nb = 3, 4
jac = vcat((hcat((Diagonal(ones(nb)) for _ in 1:nt)...) for _ in 1:nt)...) # nt×nt blocks of diagonals
colors = matrix_colors(jac)
x = ones(nt*nb)
jac2 = similar(jac)
# works without ForwardColorJacCache:
forwarddiff_color_jacobian!(jac2, (du,u)->mul!(du,jac,u), x, colorvec=colors, sparsity=jac)
# does not work with ForwardColorJacCache:
jac_cache = ForwardColorJacCache((du,u)->mul!(du,jac,u), x, colorvec=colors, sparsity=jac)
forwarddiff_color_jacobian!(jac2, (du,u)->mul!(du,jac,u), x, jac_cache)

stack trace

julia> forwarddiff_color_jacobian!(jac2, (du,u)->mul!(du,jac,u), x, jac_cache)
ERROR: MethodError: no method matching Float64(::ForwardDiff.Dual{ForwardDiff.Tag{var"#27#28", Float64}, Float64, 3})
Closest candidates are:
  (::Type{T})(::Real, ::RoundingMode) where T<:AbstractFloat at rounding.jl:200
  (::Type{T})(::T) where T<:Number at boot.jl:760
  (::Type{T})(::AbstractChar) where T<:Union{AbstractChar, Number} at char.jl:50
  ...
Stacktrace:
  [1] convert(#unused#::Type{Float64}, x::ForwardDiff.Dual{ForwardDiff.Tag{var"#27#28", Float64}, Float64, 3})
    @ Base ./number.jl:7
  [2] convert(#unused#::Type{ForwardDiff.Dual{ForwardDiff.Tag{var"#25#26", Float64}, Float64, 3}}, x::ForwardDiff.Dual{ForwardDiff.Tag{var"#27#28", Float64}, Float64, 3})
    @ ForwardDiff ~/.julia/packages/ForwardDiff/sqhTO/src/dual.jl:371
  [3] setindex!(A::Vector{ForwardDiff.Dual{ForwardDiff.Tag{var"#25#26", Float64}, Float64, 3}}, x::ForwardDiff.Dual{ForwardDiff.Tag{var"#27#28", Float64}, Float64, 3}, i1::Int64)
    @ Base ./array.jl:839
  [4] macro expansion
    @ ./broadcast.jl:984 [inlined]
  [5] macro expansion
    @ ./simdloop.jl:77 [inlined]
  [6] copyto!
    @ ./broadcast.jl:983 [inlined]
  [7] copyto!
    @ ./broadcast.jl:936 [inlined]
  [8] materialize!
    @ ./broadcast.jl:894 [inlined]
  [9] materialize!
    @ ./broadcast.jl:891 [inlined]
 [10] forwarddiff_color_jacobian!(J::SparseArrays.SparseMatrixCSC{Float64, Int64}, f::var"#27#28", x::Vector{Float64}, jac_cache::ForwardColorJacCache{Vector{ForwardDiff.Dual{ForwardDiff.Tag{var"#25#26", Float64}, Float64, 3}}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{var"#25#26", Float64}, Float64, 3}}, Vector{Float64}, Vector{Vector{Tuple{Bool, Bool, Bool}}}, Vector{Int64}, SparseArrays.SparseMatrixCSC{Float64, Int64}})
    @ SparseDiffTools ~/.julia/packages/SparseDiffTools/AhWaW/src/differentiation/compute_jacobian_ad.jl:277
 [11] top-level scope
    @ REPL[26]:1

Additional notes

The problem seems to go away when the function is not anonymous:

f!(du,u) = mul!(du,jac,u) # use f! instead of anonymous function
jac_cache = ForwardColorJacCache(f!, x, colorvec=colors, sparsity=jac)
forwarddiff_color_jacobian!(jac2, f!, x, jac_cache) # works!

But the problem comes back when the function used to build ForwardColorJacCache is a different instance than that used to call forwarddiff_color_jacobian:

g!(du,u) = mul!(du,jac,u)  # A function with a different name
jac_cache = ForwardColorJacCache(g!, x, colorvec=colors, sparsity=jac) # use g!
forwarddiff_color_jacobian!(jac2, f!, x, jac_cache) # use f! -> throws the same error

FWIW, I think I could use this second example with actually different functions that share the same sparsity for the Jacobian.