improve type stability in jacobian for a certain case (#558)

Author: KristofferC

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 #

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)

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