WIP vector_mode_dual_eval -> vector_mode_dual_eval! (#528)

Author: dlfivefifty

src/apiutils.jl

@@ -31,13 +31,13 @@ end
 
 @inline static_dual_eval(::Type{T}, f, x::StaticArray) where T = f(dualize(T, x))
 
-function vector_mode_dual_eval(f::F, x, cfg::Union{JacobianConfig,GradientConfig}) where {F}
+function vector_mode_dual_eval!(f::F, cfg::Union{JacobianConfig,GradientConfig}, x) where {F}
     xdual = cfg.duals
     seed!(xdual, x, cfg.seeds)
     return f(xdual)
 end
 
-function vector_mode_dual_eval(f!::F, y, x, cfg::JacobianConfig) where {F}
+function vector_mode_dual_eval!(f!::F, cfg::JacobianConfig, y, x) where {F}
     ydual, xdual = cfg.duals
     seed!(xdual, x, cfg.seeds)
     seed!(ydual, y)

src/gradient.jl

@@ -103,14 +103,14 @@ const GRAD_ERROR = DimensionMismatch("gradient(f, x) expects that f(x) is a real
 ###############
 
 function vector_mode_gradient(f::F, x, cfg::GradientConfig{T}) where {T, F}
-    ydual = vector_mode_dual_eval(f, x, cfg)
+    ydual = vector_mode_dual_eval!(f, cfg, x)
     ydual isa Real || throw(GRAD_ERROR)
     result = similar(x, valtype(ydual))
     return extract_gradient!(T, result, ydual)
 end
 
 function vector_mode_gradient!(result, f::F, x, cfg::GradientConfig{T}) where {T, F}
-    ydual = vector_mode_dual_eval(f, x, cfg)
+    ydual = vector_mode_dual_eval!(f, cfg, x)
     result = extract_gradient!(T, result, ydual)
     return result
 end

src/jacobian.jl

@@ -144,7 +144,7 @@ reshape_jacobian(result::DiffResult, ydual, xdual) = reshape_jacobian(DiffResult
 ###############
 
 function vector_mode_jacobian(f::F, x, cfg::JacobianConfig{T,V,N}) where {F,T,V,N}
-    ydual = vector_mode_dual_eval(f, x, cfg)
+    ydual = vector_mode_dual_eval!(f, cfg, x)
     ydual isa AbstractArray || throw(JACOBIAN_ERROR)
     result = similar(ydual, valtype(eltype(ydual)), length(ydual), N)
     extract_jacobian!(T, result, ydual, N)
@@ -153,7 +153,7 @@ function vector_mode_jacobian(f::F, x, cfg::JacobianConfig{T,V,N}) where {F,T,V,
 end
 
 function vector_mode_jacobian(f!::F, y, x, cfg::JacobianConfig{T,V,N}) where {F,T,V,N}
-    ydual = vector_mode_dual_eval(f!, y, x, cfg)
+    ydual = vector_mode_dual_eval!(f!, cfg, y, x)
     map!(d -> value(T,d), y, ydual)
     result = similar(y, length(y), N)
     extract_jacobian!(T, result, ydual, N)
@@ -162,14 +162,14 @@ function vector_mode_jacobian(f!::F, y, x, cfg::JacobianConfig{T,V,N}) where {F,
 end
 
 function vector_mode_jacobian!(result, f::F, x, cfg::JacobianConfig{T,V,N}) where {F,T,V,N}
-    ydual = vector_mode_dual_eval(f, x, 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}
-    ydual = vector_mode_dual_eval(f!, y, x, cfg)
+    ydual = vector_mode_dual_eval!(f!, cfg, y, x)
     map!(d -> value(T,d), y, ydual)
     extract_jacobian!(T, result, ydual, N)
     extract_value!(T, result, y, ydual)