From 27ad1e783772a158c7d5b7587daca8654b8da099 Mon Sep 17 00:00:00 2001 From: Albert de Montserrat <58044444+albert-de-montserrat@users.noreply.github.com> Date: Thu, 28 May 2026 17:52:10 +0200 Subject: [PATCH 1/4] Refactor vector_mode_gradient and vector_mode_jacobian Refactor gradient and jacobian functions to use direct function calls instead of static_dual_eval. --- ext/ForwardDiffStaticArraysExt.jl | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/ext/ForwardDiffStaticArraysExt.jl b/ext/ForwardDiffStaticArraysExt.jl index ec8cbc83..bf0ef99a 100644 --- a/ext/ForwardDiffStaticArraysExt.jl +++ b/ext/ForwardDiffStaticArraysExt.jl @@ -21,8 +21,6 @@ using DiffResults: DiffResult, ImmutableDiffResult, MutableDiffResult end end -@inline static_dual_eval(::Type{T}, f::F, x::StaticArray) where {T,F} = f(dualize(T, x)) - # To fix method ambiguity issues: function LinearAlgebra.eigvals(A::Symmetric{<:Dual{Tg,T,N}, <:StaticArrays.StaticMatrix}) where {Tg,T<:Real,N} return ForwardDiff._eigvals(A) @@ -54,12 +52,12 @@ end @inline function ForwardDiff.vector_mode_gradient(f::F, x::StaticArray) where {F} T = typeof(Tag(f, eltype(x))) - return extract_gradient(T, static_dual_eval(T, f, x), x) + return extract_gradient(T, f(dualize(T, x)), x) end @inline function ForwardDiff.vector_mode_gradient!(result, f::F, x::StaticArray) where {F} T = typeof(Tag(f, eltype(x))) - return extract_gradient!(T, result, static_dual_eval(T, f, x)) + return extract_gradient!(T, result, f(dualize(T, x))) end # Jacobian @@ -84,7 +82,7 @@ end @inline function ForwardDiff.vector_mode_jacobian(f::F, x::StaticArray) where {F} T = typeof(Tag(f, eltype(x))) - return extract_jacobian(T, static_dual_eval(T, f, x), x) + return extract_jacobian(T, f(dualize(T, x)), x) end function extract_jacobian(::Type{T}, ydual::AbstractArray, x::StaticArray) where T @@ -94,7 +92,7 @@ end @inline function ForwardDiff.vector_mode_jacobian!(result, f::F, x::StaticArray) where {F} T = typeof(Tag(f, eltype(x))) - ydual = static_dual_eval(T, f, x) + ydual = f(dualize(T, x)) result = extract_jacobian!(T, result, ydual, length(x)) result = extract_value!(T, result, ydual) return result @@ -102,7 +100,7 @@ end @inline function ForwardDiff.vector_mode_jacobian!(result::ImmutableDiffResult, f::F, x::StaticArray) where {F} T = typeof(Tag(f, eltype(x))) - ydual = static_dual_eval(T, f, x) + ydual = f(dualize(T, x)) result = DiffResults.jacobian!(result, extract_jacobian(T, ydual, x)) result = DiffResults.value!(Base.Fix1(value, T), result, ydual) return result From 6ede2e6e429504e21c119161da8f244ffd7dc7ee Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Thu, 4 Jun 2026 11:35:55 +0200 Subject: [PATCH 2/4] add allocation test for nested jacobians with StaticArrays --- test/AllocationsTest.jl | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/test/AllocationsTest.jl b/test/AllocationsTest.jl index 36df8eba..5321bed2 100644 --- a/test/AllocationsTest.jl +++ b/test/AllocationsTest.jl @@ -1,6 +1,7 @@ module AllocationsTest using ForwardDiff +using StaticArrays include(joinpath(dirname(@__FILE__), "utils.jl")) @@ -44,4 +45,23 @@ end @test iszero(allocs_jacobian!()) end +@testset "allocation-free nested StaticArray jacobian" begin + # test that nested jacobians of StaticArrays do not allocate. + # This is a regression test for issue #798, where the inner jacobian was allocating + @inline toy_f(x) = SVector(x[1]^2 * x[2], sin(x[1]) + x[2]^3) + + function toy_J_flat(x) + y = ForwardDiff.jacobian(toy_f, x) + return SA[y[1], y[2]] + end + + function toy_nested_jacobian(x0::SVector{2,T}) where {T} + return ForwardDiff.jacobian(toy_J_flat, x0) + end + + x0 = SVector(1.0, 2.0) + toy_nested_jacobian(x0) + @test iszero(@allocated toy_nested_jacobian(x0)) +end + end From 2efc76e08a0b6fb577a53a3865912fe840b995ae Mon Sep 17 00:00:00 2001 From: Albert de Montserrat <58044444+albert-de-montserrat@users.noreply.github.com> Date: Thu, 4 Jun 2026 16:44:35 +0200 Subject: [PATCH 3/4] Update test/AllocationsTest.jl MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: David Müller-Widmann --- test/AllocationsTest.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/AllocationsTest.jl b/test/AllocationsTest.jl index 5321bed2..ca6a05ba 100644 --- a/test/AllocationsTest.jl +++ b/test/AllocationsTest.jl @@ -48,7 +48,7 @@ end @testset "allocation-free nested StaticArray jacobian" begin # test that nested jacobians of StaticArrays do not allocate. # This is a regression test for issue #798, where the inner jacobian was allocating - @inline toy_f(x) = SVector(x[1]^2 * x[2], sin(x[1]) + x[2]^3) + toy_f(x) = SVector(x[1]^2 * x[2], sin(x[1]) + x[2]^3) function toy_J_flat(x) y = ForwardDiff.jacobian(toy_f, x) From bfc994ec299e835f7bb54975ed7cf4dc6fca2d2c Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Thu, 4 Jun 2026 16:50:58 +0200 Subject: [PATCH 4/4] remove unncessary line --- test/AllocationsTest.jl | 1 - 1 file changed, 1 deletion(-) diff --git a/test/AllocationsTest.jl b/test/AllocationsTest.jl index ca6a05ba..3858a823 100644 --- a/test/AllocationsTest.jl +++ b/test/AllocationsTest.jl @@ -60,7 +60,6 @@ end end x0 = SVector(1.0, 2.0) - toy_nested_jacobian(x0) @test iszero(@allocated toy_nested_jacobian(x0)) end