From e771e32be83bb0ceb220e246d8c5c6a4e4336b93 Mon Sep 17 00:00:00 2001
From: Kristoffer <kcarlsson89@gmail.com>
Date: Fri, 19 Nov 2021 21:53:38 +0100
Subject: [PATCH] improve type stability in jacobian for a certain case

---
 src/config.jl        |  2 +-
 src/jacobian.jl      | 12 ++++++++----
 test/JacobianTest.jl |  5 +++++
 3 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/src/config.jl b/src/config.jl
index cb23c201..30f02ae5 100644
--- a/src/config.jl
+++ b/src/config.jl
@@ -53,7 +53,7 @@ Base.copy(cfg::AbstractConfig) = deepcopy(cfg)
 
 Base.eltype(cfg::AbstractConfig) = eltype(typeof(cfg))
 
-@inline chunksize(::AbstractConfig{N}) where {N} = N
+@inline (chunksize(::AbstractConfig{N})::Int) where {N} = N
 
 ####################
 # DerivativeConfig #
diff --git a/src/jacobian.jl b/src/jacobian.jl
index 1936b36f..bcda61d7 100644
--- a/src/jacobian.jl
+++ b/src/jacobian.jl
@@ -143,7 +143,8 @@ reshape_jacobian(result::DiffResult, ydual, xdual) = reshape_jacobian(DiffResult
 # vector mode #
 ###############
 
-function vector_mode_jacobian(f::F, x, cfg::JacobianConfig{T,V,N}) where {F,T,V,N}
+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)
@@ -152,7 +153,8 @@ function vector_mode_jacobian(f::F, x, cfg::JacobianConfig{T,V,N}) where {F,T,V,
     return result
 end
 
-function vector_mode_jacobian(f!::F, y, x, cfg::JacobianConfig{T,V,N}) where {F,T,V,N}
+function vector_mode_jacobian(f!::F, y, x, cfg::JacobianConfig{T}) where {F,T}
+    N = chunksize(cfg)
     ydual = vector_mode_dual_eval!(f!, cfg, y, x)
     map!(d -> value(T,d), y, ydual)
     result = similar(y, length(y), N)
@@ -161,14 +163,16 @@ function vector_mode_jacobian(f!::F, y, x, cfg::JacobianConfig{T,V,N}) where {F,
     return result
 end
 
-function vector_mode_jacobian!(result, f::F, x, cfg::JacobianConfig{T,V,N}) where {F,T,V,N}
+function vector_mode_jacobian!(result, f::F, x, cfg::JacobianConfig{T}) where {F,T}
+    N = chunksize(cfg)
     ydual = vector_mode_dual_eval!(f, cfg, x)
     extract_jacobian!(T, result, ydual, N)
     extract_value!(T, result, ydual)
     return result
 end
 
-function vector_mode_jacobian!(result, f!::F, y, x, cfg::JacobianConfig{T,V,N}) where {F,T,V,N}
+function vector_mode_jacobian!(result, f!::F, y, x, cfg::JacobianConfig{T}) where {F,T}
+    N = chunksize(cfg)
     ydual = vector_mode_dual_eval!(f!, cfg, y, x)
     map!(d -> value(T,d), y, ydual)
     extract_jacobian!(T, result, ydual, N)
diff --git a/test/JacobianTest.jl b/test/JacobianTest.jl
index befa69ac..9ea3b0a7 100644
--- a/test/JacobianTest.jl
+++ b/test/JacobianTest.jl
@@ -237,4 +237,9 @@ end
     @test ForwardDiff.jacobian(x -> eigvals(Symmetric(x*x')), [1.,2.]) ≈ [0 0; 2 4]
 end
 
+@testset "type stability" begin
+    g!(dy, y) = dy[1] = y[1]
+    @inferred ForwardDiff.jacobian(g!, [1.0], [0.0])
+end
+
 end # module