From 47adb2c842fa8cdd3dd4b2c72af0a6f0a70eb997 Mon Sep 17 00:00:00 2001 From: David Widmann Date: Wed, 3 Sep 2025 23:03:52 +0200 Subject: [PATCH 1/2] Fix gradient and Jacobian for functions with `Dual` output --- ext/ForwardDiffStaticArraysExt.jl | 6 +++--- src/dual.jl | 11 +++++++++++ src/gradient.jl | 4 ++-- src/jacobian.jl | 4 ++-- test/DualTest.jl | 15 +++++++++++++++ test/GradientTest.jl | 31 +++++++++++++++++++++++++++++++ test/JacobianTest.jl | 31 +++++++++++++++++++++++++++++++ 7 files changed, 95 insertions(+), 7 deletions(-) diff --git a/ext/ForwardDiffStaticArraysExt.jl b/ext/ForwardDiffStaticArraysExt.jl index e97da9d5..ec8cbc83 100644 --- 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 @@ -77,7 +77,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 @@ -88,7 +88,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 index 69a36760..f9b83c4e 100644 --- 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 index 98bdb091..e6e56e46 100644 --- 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 index d654d38b..83e21ea5 100644 --- 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 index b64094ea..fa26bfd3 100644 --- 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 index 4f46c167..82c2b2a8 100644 --- 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 index b010078f..9cc5024c 100644 --- 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 From dcde31e7769bbf34a31eb65df69aaab00f1b71a6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Tue, 9 Sep 2025 15:40:08 +0000 Subject: [PATCH 2/2] Bump version from 1.2.0 to 1.2.1 --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 24885b11..f631c512 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "ForwardDiff" uuid = "f6369f11-7733-5829-9624-2563aa707210" -version = "1.2.0" +version = "1.2.1" [deps] CommonSubexpressions = "bbf7d656-a473-5ed7-a52c-81e309532950"