specialize on input function on a few more wrapper functions

Author: KristofferC

src/apiutils.jl

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

src/gradient.jl

@@ -29,7 +29,7 @@ Compute `∇f` evaluated at `x` and store the result(s) in `result`, assuming `f
 This method assumes that `isa(f(x), Real)`.
 
 """
-function gradient!(result::Union{AbstractArray,DiffResult}, f, x::AbstractArray, cfg::GradientConfig{T} = GradientConfig(f, x), ::Val{CHK}=Val{true}()) where {T, CHK}
+function gradient!(result::Union{AbstractArray,DiffResult}, f::F, x::AbstractArray, cfg::GradientConfig{T} = GradientConfig(f, x), ::Val{CHK}=Val{true}()) where {T, CHK, F}
     CHK && checktag(T, f, x)
     if chunksize(cfg) == length(x)
         vector_mode_gradient!(result, f, x, cfg)
@@ -92,13 +92,13 @@ end
 # vector mode #
 ###############
 
-function vector_mode_gradient(f, x, cfg::GradientConfig{T}) where {T}
+function vector_mode_gradient(f::F, x, cfg::GradientConfig{T}) where {T, F}
     ydual = vector_mode_dual_eval(f, x, cfg)
     result = similar(x, valtype(ydual))
     return extract_gradient!(T, result, ydual)
 end
 
-function vector_mode_gradient!(result, f, x, cfg::GradientConfig{T}) where {T}
+function vector_mode_gradient!(result, f::F, x, cfg::GradientConfig{T}) where {T, F}
     ydual = vector_mode_dual_eval(f, x, cfg)
     result = extract_gradient!(T, result, ydual)
     return result