Fix gradient and Jacobian for functions with Dual output

devmotion · devmotion · commit aa14fbe67d2d · 2025-09-03T23:14:30.000+02:00
diff --git a/ext/ForwardDiffStaticArraysExt.jl b/ext/ForwardDiffStaticArraysExt.jl
@@ -47,7 +47,7 @@ ForwardDiff._lyap_div!!(A::StaticArrays.MMatrix, λ::AbstractVector) = ForwardDi
     result = Expr(:tuple, [:(partials(T, y, $i)) for i in 1:length(x)]...)
     return quote
         $(Expr(:meta, :inline))
-        V = StaticArrays.similar_type(S, valtype($y))
+        V = StaticArrays.similar_type(S, valtype(T, $y))
         return V($result)
     end
 end
@@ -76,7 +76,7 @@ end
     result = Expr(:tuple, [:(partials(T, ydual[$i], $j)) for i in 1:M, j in 1:N]...)
     return quote
         $(Expr(:meta, :inline))
-        V = StaticArrays.similar_type(S, valtype(eltype($ydual)), Size($M, $N))
+        V = StaticArrays.similar_type(S, valtype(T, eltype($ydual)), Size($M, $N))
         return V($result)
     end
 end
@@ -87,7 +87,7 @@ end
 end
 
 function extract_jacobian(::Type{T}, ydual::AbstractArray, x::StaticArray) where T
-    result = similar(ydual, valtype(eltype(ydual)), length(ydual), length(x))
+    result = similar(ydual, valtype(T, eltype(ydual)), length(ydual), length(x))
     return extract_jacobian!(T, result, ydual, length(x))
 end
 
diff --git a/src/dual.jl b/src/dual.jl
@@ -128,6 +128,17 @@ end
 @inline valtype(::Dual{T,V,N}) where {T,V,N} = V
 @inline valtype(::Type{Dual{T,V,N}}) where {T,V,N} = V
 
+@inline valtype(::Type{T}, ::V) where {T,V} = valtype(T, V)
+@inline valtype(::Type, ::Type{V}) where {V} = V
+@inline valtype(::Type{T}, ::Type{Dual{T,V,N}}) where {T,V,N} = V
+@inline function valtype(::Type{T}, ::Type{Dual{S,V,N}}) where {T,S,V,N}
+    if S ≺ T
+        Dual{S,V,N}
+    else
+        throw(DualMismatchError(T,S))
+    end
+end
+
 @inline tagtype(::V) where {V} = Nothing
 @inline tagtype(::Type{V}) where {V} = Nothing
 @inline tagtype(::Dual{T,V,N}) where {T,V,N} = T
diff --git a/src/gradient.jl b/src/gradient.jl
@@ -97,7 +97,7 @@ const GRAD_ERROR = DimensionMismatch("gradient(f, x) expects that f(x) is a real
 function vector_mode_gradient(f::F, x, cfg::GradientConfig{T}) where {T, F}
     ydual = vector_mode_dual_eval!(f, cfg, x)
     ydual isa Real || throw(GRAD_ERROR)
-    result = similar(x, valtype(ydual))
+    result = similar(x, valtype(T, ydual))
     return extract_gradient!(T, result, ydual)
 end
 
@@ -156,7 +156,7 @@ function chunk_mode_gradient_expr(result_definition::Expr)
 end
 
 @eval function chunk_mode_gradient(f::F, x, cfg::GradientConfig{T,V,N}) where {F,T,V,N}
-    $(chunk_mode_gradient_expr(:(result = similar(x, valtype(ydual)))))
+    $(chunk_mode_gradient_expr(:(result = similar(x, valtype(T, ydual)))))
 end
 
 @eval function chunk_mode_gradient!(result, f::F, x, cfg::GradientConfig{T,V,N}) where {F,T,V,N}
diff --git a/src/jacobian.jl b/src/jacobian.jl
@@ -128,7 +128,7 @@ function vector_mode_jacobian(f::F, x, cfg::JacobianConfig{T}) where {F,T}
     N = chunksize(cfg)
     ydual = vector_mode_dual_eval!(f, cfg, x)
     ydual isa AbstractArray || throw(JACOBIAN_ERROR)
-    result = similar(ydual, valtype(eltype(ydual)), length(ydual), N)
+    result = similar(ydual, valtype(T, eltype(ydual)), length(ydual), N)
     extract_jacobian!(T, result, ydual, N)
     extract_value!(T, result, ydual)
     return result
@@ -217,7 +217,7 @@ end
                                    seed!(xdual, x)
                                end,
                                :(ydual = f(xdual)),
-                               :(result = similar(ydual, valtype(eltype(ydual)), length(ydual), xlen)),
+                               :(result = similar(ydual, valtype(T, eltype(ydual)), length(ydual), xlen)),
                                :()))
 end
 
diff --git a/test/DualTest.jl b/test/DualTest.jl
@@ -101,6 +101,21 @@ ForwardDiff.:≺(::Type{OuterTestTag}, ::Type{TestTag}) = false
     @test ForwardDiff.valtype(NESTED_FDNUM) == Dual{TestTag,V,M}
     @test ForwardDiff.valtype(typeof(NESTED_FDNUM)) == Dual{TestTag,V,M}
 
+    @test ForwardDiff.valtype(TestTag, FDNUM) == V
+    @test ForwardDiff.valtype(TestTag, typeof(FDNUM)) == V
+    @test ForwardDiff.valtype(TestTag, NESTED_FDNUM) == Dual{TestTag,V,M}
+    @test ForwardDiff.valtype(TestTag, typeof(NESTED_FDNUM)) == Dual{TestTag,V,M}
+
+    @test ForwardDiff.valtype(OuterTestTag, FDNUM) == Dual{TestTag,V,N}
+    @test ForwardDiff.valtype(OuterTestTag, typeof(FDNUM)) == Dual{TestTag,V,N}
+    @test ForwardDiff.valtype(OuterTestTag, NESTED_FDNUM) == Dual{TestTag,Dual{TestTag,V,M},N}
+    @test ForwardDiff.valtype(OuterTestTag, typeof(NESTED_FDNUM)) == Dual{TestTag,Dual{TestTag,V,M},N}
+
+    @test_throws ForwardDiff.DualMismatchError(TestTag, OuterTestTag) ForwardDiff.valtype(TestTag, Dual{OuterTestTag}(PRIMAL, PARTIALS))
+    @test_throws ForwardDiff.DualMismatchError(TestTag, OuterTestTag) ForwardDiff.valtype(TestTag, typeof(Dual{OuterTestTag}(PRIMAL, PARTIALS)))
+    @test_throws ForwardDiff.DualMismatchError(TestTag, OuterTestTag) ForwardDiff.valtype(TestTag, Dual{OuterTestTag}(Dual{TestTag}(PRIMAL, M_PARTIALS), NESTED_PARTIALS))
+    @test_throws ForwardDiff.DualMismatchError(TestTag, OuterTestTag) ForwardDiff.valtype(TestTag, typeof(Dual{OuterTestTag}(Dual{TestTag}(PRIMAL, M_PARTIALS), NESTED_PARTIALS)))
+
     #####################
     # Generic Functions #
     #####################
diff --git a/test/GradientTest.jl b/test/GradientTest.jl
@@ -12,6 +12,11 @@ using DiffTests
 
 include(joinpath(dirname(@__FILE__), "utils.jl"))
 
+struct TestTag end
+struct OuterTestTag end
+ForwardDiff.:≺(::Type{TestTag}, ::Type{OuterTestTag}) = true
+ForwardDiff.:≺(::Type{OuterTestTag}, ::Type{<:Tag}) = true
+
 ##################
 # hardcoded test #
 ##################
@@ -255,4 +260,30 @@ end
     end
 end
 
+# issue #769
+@testset "functions with `Dual` output" begin
+    x = [Dual{OuterTestTag}(Dual{TestTag}(1.3, 2.1), Dual{TestTag}(0.3, -2.4))]
+    f(x) = sum(ForwardDiff.value, x)
+    der = ForwardDiff.derivative(ForwardDiff.value, only(x))
+
+    # Vector mode
+    grad = ForwardDiff.gradient(f, x)
+    @test grad isa Vector{typeof(der)}
+    @test grad == [der]
+    grad = ForwardDiff.gradient(f, SVector{1}(x))
+    @test grad isa SVector{1,typeof(der)}
+    @test grad == SVector{1}(der)
+
+    # Chunk mode
+    y = repeat(x, 3)
+    cfg = ForwardDiff.GradientConfig(f, y, ForwardDiff.Chunk{2}())
+    grad = ForwardDiff.gradient(f, y, cfg)
+    @test grad isa Vector{typeof(der)}
+    @test grad == [der, der, der]
+    cfg = ForwardDiff.GradientConfig(f, SVector{3}(y), ForwardDiff.Chunk{2}())
+    grad = ForwardDiff.gradient(f, SVector{3}(y), cfg)
+    @test grad isa SVector{3,typeof(der)}
+    @test grad == SVector{3}(der, der, der)
+end
+
 end # module
diff --git a/test/JacobianTest.jl b/test/JacobianTest.jl
@@ -11,6 +11,11 @@ using LinearAlgebra
 
 include(joinpath(dirname(@__FILE__), "utils.jl"))
 
+struct TestTag end
+struct OuterTestTag end
+ForwardDiff.:≺(::Type{TestTag}, ::Type{OuterTestTag}) = true
+ForwardDiff.:≺(::Type{OuterTestTag}, ::Type{<:Tag}) = true
+
 ##################
 # hardcoded test #
 ##################
@@ -308,4 +313,30 @@ end
     end
 end
 
+# issue #769
+@testset "functions with `Dual` output" begin
+    x = [Dual{OuterTestTag}(Dual{TestTag}(1.3, 2.1), Dual{TestTag}(0.3, -2.4))]
+    f(x) = map(ForwardDiff.value, x)
+    der = ForwardDiff.derivative(ForwardDiff.value, only(x))
+
+    # Vector mode
+    jac = ForwardDiff.jacobian(f, x)
+    @test jac isa Matrix{typeof(der)}
+    @test jac == [der;;]
+    jac = ForwardDiff.jacobian(f, SVector{1}(x))
+    @test jac isa SMatrix{1,1,typeof(der)}
+    @test jac == SMatrix{1,1}(der)
+
+    # Chunk mode
+    y = repeat(x, 3)
+    cfg = ForwardDiff.JacobianConfig(f, y, ForwardDiff.Chunk{2}())
+    jac = ForwardDiff.jacobian(f, y, cfg)
+    @test jac isa Matrix{typeof(der)}
+    @test jac == Diagonal([der, der, der])
+    cfg = ForwardDiff.JacobianConfig(f, SVector{3}(y), ForwardDiff.Chunk{2}())
+    jac = ForwardDiff.jacobian(f, SVector{3}(y), cfg)
+    @test jac isa SMatrix{3,3,typeof(der)}
+    @test jac == Diagonal([der, der, der])
+end
+
 end # module
