From fdcaadbaee09d7690857217bda2164c07c0e04cf Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A9s=20Riedemann?=
 <38795484+longemen3000@users.noreply.github.com>
Date: Tue, 13 May 2025 22:55:48 -0400
Subject: [PATCH 1/5] improve performance of hessians with static arrays

---
 ext/ForwardDiffStaticArraysExt.jl | 10 ++++++++++
 test/HessianTest.jl               |  8 ++++++++
 2 files changed, 18 insertions(+)

diff --git a/ext/ForwardDiffStaticArraysExt.jl b/ext/ForwardDiffStaticArraysExt.jl
index 63f841db..8226aea3 100644
--- a/ext/ForwardDiffStaticArraysExt.jl
+++ b/ext/ForwardDiffStaticArraysExt.jl
@@ -81,6 +81,16 @@ end
     end
 end
 
+@generated function extract_jacobian(::Type{T}, ydual::Partials{M}, x::S) where {M, T, S<:StaticArray}
+    N = length(x)
+    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))
+        return V($result)
+    end
+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)
diff --git a/test/HessianTest.jl b/test/HessianTest.jl
index 4c667e5e..b826c2ff 100644
--- a/test/HessianTest.jl
+++ b/test/HessianTest.jl
@@ -163,4 +163,12 @@ end
     @test ForwardDiff.hessian(x->dot(x,H,x), zeros(3)) ≈ [2 6 10; 6 10 14; 10 14 18]
 end
 
+@testset "allocation-free hessian with StaticArrays" begin
+    #https://github.com/JuliaDiff/ForwardDiff.jl/issues/720
+    g = r -> (r[1]^2 - 3) * (r[2]^2 - 2)
+    x = SA_F32[0.5, 2.7]
+    hres = DiffResults.HessianResult(x)
+    @test @allocated(ForwardDiff.hessian!(hres, g, x))) == 0
+end
+
 end # module

From 183e29033a06e95cc35a47c8ead3ca96bbe4c59a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A9s=20Riedemann?=
 <38795484+longemen3000@users.noreply.github.com>
Date: Tue, 13 May 2025 23:02:53 -0400
Subject: [PATCH 2/5] ForwardDiffStaticArraysExt.jl: import Partials

---
 ext/ForwardDiffStaticArraysExt.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ext/ForwardDiffStaticArraysExt.jl b/ext/ForwardDiffStaticArraysExt.jl
index 8226aea3..fe65eb33 100644
--- a/ext/ForwardDiffStaticArraysExt.jl
+++ b/ext/ForwardDiffStaticArraysExt.jl
@@ -3,7 +3,7 @@ module ForwardDiffStaticArraysExt
 using ForwardDiff, StaticArrays
 using ForwardDiff.LinearAlgebra
 using ForwardDiff.DiffResults
-using ForwardDiff: Dual, partials, GradientConfig, JacobianConfig, HessianConfig, Tag, Chunk,
+using ForwardDiff: Dual, partials, Partials, GradientConfig, JacobianConfig, HessianConfig, Tag, Chunk,
                    gradient, hessian, jacobian, gradient!, hessian!, jacobian!,
                    extract_gradient!, extract_jacobian!, extract_value!,
                    vector_mode_gradient, vector_mode_gradient!,

From bade2d24298afd9e19a210ea706e3c23e53de2f3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A9s=20Riedemann?=
 <38795484+longemen3000@users.noreply.github.com>
Date: Tue, 13 May 2025 23:11:11 -0400
Subject: [PATCH 3/5] Update HessianTest.jl

---
 test/HessianTest.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/HessianTest.jl b/test/HessianTest.jl
index b826c2ff..051377fe 100644
--- a/test/HessianTest.jl
+++ b/test/HessianTest.jl
@@ -168,7 +168,7 @@ end
     g = r -> (r[1]^2 - 3) * (r[2]^2 - 2)
     x = SA_F32[0.5, 2.7]
     hres = DiffResults.HessianResult(x)
-    @test @allocated(ForwardDiff.hessian!(hres, g, x))) == 0
+    @test @allocated(ForwardDiff.hessian!(hres, g, x)) == 0
 end
 
 end # module

From 190e85ffe09fa6d07e711a27922cbefa455cd008 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A9s=20Riedemann?=
 <38795484+longemen3000@users.noreply.github.com>
Date: Tue, 13 May 2025 23:21:11 -0400
Subject: [PATCH 4/5] Update HessianTest.jl

---
 test/HessianTest.jl | 1 +
 1 file changed, 1 insertion(+)

diff --git a/test/HessianTest.jl b/test/HessianTest.jl
index 051377fe..9f12f91c 100644
--- a/test/HessianTest.jl
+++ b/test/HessianTest.jl
@@ -168,6 +168,7 @@ end
     g = r -> (r[1]^2 - 3) * (r[2]^2 - 2)
     x = SA_F32[0.5, 2.7]
     hres = DiffResults.HessianResult(x)
+    ForwardDiff.hessian!(hres, g, x)
     @test @allocated(ForwardDiff.hessian!(hres, g, x)) == 0
 end
 

From a6306bfb4cd339164214d43cfd125faccfb12d45 Mon Sep 17 00:00:00 2001
From: David Widmann <devmotion@users.noreply.github.com>
Date: Tue, 9 Sep 2025 12:04:29 +0200
Subject: [PATCH 5/5] Simplify code and fix test

---
 ext/ForwardDiffStaticArraysExt.jl | 15 +++------------
 test/HessianTest.jl               | 14 ++++++++------
 2 files changed, 11 insertions(+), 18 deletions(-)

diff --git a/ext/ForwardDiffStaticArraysExt.jl b/ext/ForwardDiffStaticArraysExt.jl
index fe65eb33..e97da9d5 100644
--- a/ext/ForwardDiffStaticArraysExt.jl
+++ b/ext/ForwardDiffStaticArraysExt.jl
@@ -3,7 +3,7 @@ module ForwardDiffStaticArraysExt
 using ForwardDiff, StaticArrays
 using ForwardDiff.LinearAlgebra
 using ForwardDiff.DiffResults
-using ForwardDiff: Dual, partials, Partials, GradientConfig, JacobianConfig, HessianConfig, Tag, Chunk,
+using ForwardDiff: Dual, partials, npartials, Partials, GradientConfig, JacobianConfig, HessianConfig, Tag, Chunk,
                    gradient, hessian, jacobian, gradient!, hessian!, jacobian!,
                    extract_gradient!, extract_jacobian!, extract_value!,
                    vector_mode_gradient, vector_mode_gradient!,
@@ -71,17 +71,8 @@ end
 @inline ForwardDiff.jacobian!(result::Union{AbstractArray,DiffResult}, f::F, x::StaticArray, cfg::JacobianConfig) where {F} = jacobian!(result, f, x)
 @inline ForwardDiff.jacobian!(result::Union{AbstractArray,DiffResult}, f::F, x::StaticArray, cfg::JacobianConfig, ::Val) where {F} = jacobian!(result, f, x)
 
-@generated function extract_jacobian(::Type{T}, ydual::StaticArray, x::S) where {T,S<:StaticArray}
-    M, N = length(ydual), length(x)
-    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))
-        return V($result)
-    end
-end
-
-@generated function extract_jacobian(::Type{T}, ydual::Partials{M}, x::S) where {M, T, S<:StaticArray}
+@generated function extract_jacobian(::Type{T}, ydual::Union{StaticArray,Partials}, x::S) where {T,S<:StaticArray}
+    M = ydual <: Partials ? npartials(ydual) : length(ydual)
     N = length(x)
     result = Expr(:tuple, [:(partials(T, ydual[$i], $j)) for i in 1:M, j in 1:N]...)
     return quote
diff --git a/test/HessianTest.jl b/test/HessianTest.jl
index 9f12f91c..8be72ee5 100644
--- a/test/HessianTest.jl
+++ b/test/HessianTest.jl
@@ -163,13 +163,15 @@ end
     @test ForwardDiff.hessian(x->dot(x,H,x), zeros(3)) ≈ [2 6 10; 6 10 14; 10 14 18]
 end
 
+#https://github.com/JuliaDiff/ForwardDiff.jl/issues/720
 @testset "allocation-free hessian with StaticArrays" begin
-    #https://github.com/JuliaDiff/ForwardDiff.jl/issues/720
-    g = r -> (r[1]^2 - 3) * (r[2]^2 - 2)
-    x = SA_F32[0.5, 2.7]
-    hres = DiffResults.HessianResult(x)
-    ForwardDiff.hessian!(hres, g, x)
-    @test @allocated(ForwardDiff.hessian!(hres, g, x)) == 0
+    function hessian_allocs()
+        g = r -> (r[1]^2 - 3) * (r[2]^2 - 2)
+        x = SVector(0.5, 2.8)
+        hres = DiffResults.HessianResult(x)
+        return @allocated(ForwardDiff.hessian!(hres, g, x))
+    end
+    @test iszero(hessian_allocs())
 end
 
 end # module