Problem with nothing tag in Autodiff.jl for mixed spatial and DOF based derivatives

[kishore-nori]

Currently Gridap at src/Arrays/Autodiff.jl uses nothing based ForwardDiff.Tag for constructing the configs for ForwardDiff.gradient and ForwardDiff.jacobian,

Gridap.jl/src/Arrays/Autodiff.jl

Lines 51 to 59 in 78b27c4

	function return_cache(k::ConfigMap{typeof(ForwardDiff.gradient)},x)
	cfg = ForwardDiff.GradientConfig(nothing,x,ForwardDiff.Chunk{length(x)}())
	cfg
	end
	function return_cache(k::ConfigMap{typeof(ForwardDiff.jacobian)},x)
	cfg = ForwardDiff.JacobianConfig(nothing,x,ForwardDiff.Chunk{length(x)}())
	cfg
	end

this is not a problem for spatial gradient/laplacians of FEFunctions, as ForwardDiff machinery is not being used inside. Currently, I am working with CellFields where the spatial gradient/laplacian requires the usage of ForwardDiff and the interaction of the tag generated in this ForwardDiff usage has interaction problem with nothing tag from derivatives with respect DOF with respect to the functionals involving the spatial derivative, related to determining the order of the tags. Following is the MWE with only ForwardDiff:

using ForwardDiff
using LinearAlgebra

f(θ,x) = θ⋅x

function test_func(θ)
  _f(x) = f(θ,x)
  x = [0.5,0.6]
  norm(ForwardDiff.gradient(_f,x))
end

θ = [2.0,0.5]
# following the strategy of `Gridap.Arrays` `Autodiff.jl`
cfg = ForwardDiff.GradientConfig(nothing,θ,ForwardDiff.Chunk{length(θ)}())
cfg isa ForwardDiff.GradientConfig{Nothing}
function cfg_type(cfg::ForwardDiff.GradientConfig{T}) where T
  T
end
cfg_type(cfg)

xdual = cfg.duals
ForwardDiff.seed!(xdual, θ, cfg.seeds)

ydual = test_func(xdual) # this is the place where tag order error occurs 
result = similar(θ, ForwardDiff.valtype(ydual))

result = ForwardDiff.extract_gradient!(cfg_type(cfg),result,ydual)

ydual = test_func(xdual) call results in the following error, this where the interaction of dual numbers from different ForwardDiff calls occurs. This is due to one tag being Nothing.

ERROR: Cannot determine ordering of Dual tags ForwardDiff.Tag{var"#_f#12"{Vector{ForwardDiff.Dual{Nothing, Float64, 2}}}, Float64} and Nothing
Stacktrace:
  [1] ≺(a::Type, b::Type)
    @ ForwardDiff ~/.julia/packages/ForwardDiff/wAaVJ/src/dual.jl:54
  [2] *
    @ ~/.julia/packages/ForwardDiff/wAaVJ/src/dual.jl:140 [inlined]
  [3] dot(x::ForwardDiff.Dual{Nothing, Float64, 2}, y::ForwardDiff.Dual{ForwardDiff.Tag{var"#_f#12"{Vector{ForwardDiff.Dual{Nothing, Float64, 2}}}, Float64}, Float64, 2})
    @ LinearAlgebra ~/julia-src/julia-1.7.2/share/julia/stdlib/v1.7/LinearAlgebra/src/generic.jl:905
  [4] dot(x::Vector{ForwardDiff.Dual{Nothing, Float64, 2}}, y::Vector{ForwardDiff.Dual{ForwardDiff.Tag{var"#_f#12"{Vector{ForwardDiff.Dual{Nothing, Float64, 2}}}, Float64}, Float64, 2}})
    @ LinearAlgebra ~/julia-src/julia-1.7.2/share/julia/stdlib/v1.7/LinearAlgebra/src/generic.jl:915
  [5] f(θ::Vector{ForwardDiff.Dual{Nothing, Float64, 2}}, x::Vector{ForwardDiff.Dual{ForwardDiff.Tag{var"#_f#12"{Vector{ForwardDiff.Dual{Nothing, Float64, 2}}}, Float64}, Float64, 2}})
    @ Main ~/Documents/Special/PhD-learning-work-notes/Work/Gridap-Ecosystem/MWEsGridap/ForwardDiff-nothing-tag-order-mwe.jl:4
  [6] (::var"#_f#12"{Vector{ForwardDiff.Dual{Nothing, Float64, 2}}})(x::Vector{ForwardDiff.Dual{ForwardDiff.Tag{var"#_f#12"{Vector{ForwardDiff.Dual{Nothing, Float64, 2}}}, Float64}, Float64, 2}})
    @ Main ~/Documents/Special/PhD-learning-work-notes/Work/Gridap-Ecosystem/MWEsGridap/ForwardDiff-nothing-tag-order-mwe.jl:7
  [7] vector_mode_dual_eval!(f::var"#_f#12"{Vector{ForwardDiff.Dual{Nothing, Float64, 2}}}, cfg::ForwardDiff.GradientConfig{ForwardDiff.Tag{var"#_f#12"{Vector{ForwardDiff.Dual{Nothing, Float64, 2}}}, Float64}, Float64, 2, Vector{ForwardDiff.Dual{ForwardDiff.Tag{var"#_f#12"{Vector{ForwardDiff.Dual{Nothing, Float64, 2}}}, Float64}, Float64, 2}}}, x::Vector{Float64})
    @ ForwardDiff ~/.julia/packages/ForwardDiff/wAaVJ/src/apiutils.jl:37
  [8] vector_mode_gradient(f::var"#_f#12"{Vector{ForwardDiff.Dual{Nothing, Float64, 2}}}, x::Vector{Float64}, cfg::ForwardDiff.GradientConfig{ForwardDiff.Tag{var"#_f#12"{Vector{ForwardDiff.Dual{Nothing, Float64, 2}}}, Float64}, Float64, 2, Vector{ForwardDiff.Dual{ForwardDiff.Tag{var"#_f#12"{Vector{ForwardDiff.Dual{Nothing, Float64, 2}}}, Float64}, Float64, 2}}})
    @ ForwardDiff ~/.julia/packages/ForwardDiff/wAaVJ/src/gradient.jl:106
  [9] gradient(f::Function, x::Vector{Float64}, cfg::ForwardDiff.GradientConfig{ForwardDiff.Tag{var"#_f#12"{Vector{ForwardDiff.Dual{Nothing, Float64, 2}}}, Float64}, Float64, 2, Vector{ForwardDiff.Dual{ForwardDiff.Tag{var"#_f#12"{Vector{ForwardDiff.Dual{Nothing, Float64, 2}}}, Float64}, Float64, 2}}}, ::Val{true})
    @ ForwardDiff ~/.julia/packages/ForwardDiff/wAaVJ/src/gradient.jl:19
 [10] gradient(f::Function, x::Vector{Float64}, cfg::ForwardDiff.GradientConfig{ForwardDiff.Tag{var"#_f#12"{Vector{ForwardDiff.Dual{Nothing, Float64, 2}}}, Float64}, Float64, 2, Vector{ForwardDiff.Dual{ForwardDiff.Tag{var"#_f#12"{Vector{ForwardDiff.Dual{Nothing, Float64, 2}}}, Float64}, Float64, 2}}}) (repeats 2 times)
    @ ForwardDiff ~/.julia/packages/ForwardDiff/wAaVJ/src/gradient.jl:17
 [11] test_func(θ::Vector{ForwardDiff.Dual{Nothing, Float64, 2}})

cc @amartinhuertas @fverdugo @santiagobadia

The possible fix is to have a dummy tag for AD functionality for functionals in Gridap, instead of nothing. Having the following works without problems and the results match with manual gradients and FiniteDifferences for complicated cases:

dummy() = nothing # a better name can be chosen 

function return_cache(k::ConfigMap{typeof(ForwardDiff.gradient)},x) 
   cfg = ForwardDiff.GradientConfig(dummy,x,ForwardDiff.Chunk{length(x)}()) 
   cfg 
 end 
  
 function return_cache(k::ConfigMap{typeof(ForwardDiff.jacobian)},x) 
   cfg = ForwardDiff.JacobianConfig(dummy,x,ForwardDiff.Chunk{length(x)}()) 
   cfg 
 end 

Introducing this would not disturb or break anything or fail in the existing use cases.

[amartinhuertas]

For the records, we talked about having a more general approach to solve this issue. Instead of a fixed dummy function, use the function at hand that we are about to differentiate. To this end we need to modify the ConfigMap struct as

struct ConfigMap{F} <: Map
 f::F
end
ConfigMap()=ConfigMap(nothing) # required for backward compatibility

[kishore-nori]

The above approach seems to make things slower, please check #806 (comment). I feel we need an extra input parameter in gradient, jacobian (or _gradient or _jacobian so as to not disturb the high-level API) when we want to trigger tag checking, this is required only when we have CellFields where ForwardDiff will be use and not in FEFunction cases. See the section in ForwardDiff.jl docs on tag checking. We can use a input parameter like Val{true/false}.

[kishore-nori]

Issue resolved in PR #806 and merged with master branch.