improve performance of hessians with static arrays

Author: longemen3000

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)

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