ForwardDiff tagcount with custom diff rules

Specific Domains Numerics
,
1

I’m writing a custom gradient for a function that takes two derivatives using ForwardDiff. I receive an error unless I specify the tag for the returned Dual. I figured that the tag I should return depends on the tag count, but it’s not clear to me whether I should attach the tag with the larger or the smaller count. Here’s an implementation that seems to work, though I’m not sure if it’s correct.

function JacobiElliptic.CarlsonAlg.E(x::ForwardDiff.Dual{T1}, y::ForwardDiff.Dual{T2}) where {T1, T2}
    xval = x.value
    yval = y.value

    ∂yE = iszero(yval) ? -π/8 : (JacobiElliptic.CarlsonAlg.E(xval, yval)-JacobiElliptic.CarlsonAlg.F(xval, yval))/(2yval)
    return ForwardDiff.Dual{ForwardDiff.tagcount(T1) < ForwardDiff.tagcount(T2) ? T1 : T2}(JacobiElliptic.CarlsonAlg.E(xval, yval), ForwardDiff.Partials(((sqrt(1-yval*sin(xval)^2)), ∂yE)))
end

Which is the correct tag that I should attach to this method?

2

Hi! Can you share a complete example with the error you get from ForwardDiff?

3

I get the following error if I return a Dual without specifying the tag:

function JacobiElliptic.CarlsonAlg.E(x::ForwardDiff.Dual{T1}, y::ForwardDiff.Dual{T2}) where {T1, T2}
    xval = x.value
    yval = y.value

    ∂yE = iszero(yval) ? -π/8 : (JacobiElliptic.CarlsonAlg.E(xval, yval)-JacobiElliptic.CarlsonAlg.F(xval, yval))/(2yval)
    ForwardDiff.Dual(JacobiElliptic.CarlsonAlg.E(xval, yval), ForwardDiff.Partials(((sqrt(1-yval*sin(xval)^2)), ∂yE)))
end

ERROR: Cannot determine ordering of Dual tags Nothing and ForwardDiff.Tag{typeof(func), Float64}
Stacktrace:
 [1] ≺(a::Type, b::Type)
   @ ForwardDiff ~/.julia/packages/ForwardDiff/PcZ48/src/dual.jl:54
 [2] partials
   @ ~/.julia/packages/ForwardDiff/PcZ48/src/dual.jl:108 [inlined]
 [3] extract_gradient!(::Type{ForwardDiff.Tag{…}}, result::Vector{Float64}, dual::Dual{Nothing, Float64, 2})
   @ ForwardDiff ~/.julia/packages/ForwardDiff/PcZ48/src/gradient.jl:64
 [4] vector_mode_gradient(f::typeof(func), x::Vector{…}, cfg::ForwardDiff.GradientConfig{…})
   @ ForwardDiff ~/.julia/packages/ForwardDiff/PcZ48/src/gradient.jl:92
 [5] gradient(f::Function, x::Vector{…}, cfg::ForwardDiff.GradientConfig{…}, ::Val{…})
   @ ForwardDiff ~/.julia/packages/ForwardDiff/PcZ48/src/gradient.jl:0
 [6] gradient(f::Function, x::Vector{Float64}, cfg::ForwardDiff.GradientConfig{ForwardDiff.Tag{…}, Float64, 2, Vector{…}})
   @ ForwardDiff ~/.julia/packages/ForwardDiff/PcZ48/src/gradient.jl:17

My main confusion is that I’m not sure which tag to return since I’ve defined the function to act on two separate Duals.

4

I think you probably want this:

function JacobiElliptic.CarlsonAlg.E(x::ForwardDiff.Dual{T}, y::ForwardDiff.Dual{T}) where {T}
    xval = x.value
    yval = y.value

    ∂yE = iszero(yval) ? -π/8 : (JacobiElliptic.CarlsonAlg.E(xval, yval)-JacobiElliptic.CarlsonAlg.F(xval, yval))/(2yval)
    ForwardDiff.Dual{T}(JacobiElliptic.CarlsonAlg.E(xval, yval), ForwardDiff.Partials(((sqrt(1-yval*sin(xval)^2)), ∂yE)))
end

Two different tags means these are independent perturbations, which will occur when taking second derivatives – and will result in nested Dual types:

julia> using ForwardDiff: Dual, Tag, derivative

julia> Dual(1,2) + Dual(3,5)  # first derivative, both arguments perturbed together
Dual{Nothing}(4,7)

julia> Dual{Tag{1,1}}(1,2) + Dual{Tag{2,2}}(3,5)  # second derivative, two indep perturbations
Dual{Tag{2, 2}}(Dual{Tag{1, 1}}(4,2),Dual{Tag{1, 1}}(5,0))

julia> Dual{Tag{2,2}}(1,2) + Dual{Tag{1,1}}(3,5)  # same, but organised differently
Dual{Tag{2, 2}}(Dual{Tag{1, 1}}(4,5),Dual{Tag{1, 1}}(2,0))
1 Like
5

Thanks for that