diff --git a/ext/ForwardDiffStaticArraysExt.jl b/ext/ForwardDiffStaticArraysExt.jl index cb9222ff..1ec87ea7 100644 --- a/ext/ForwardDiffStaticArraysExt.jl +++ b/ext/ForwardDiffStaticArraysExt.jl @@ -37,13 +37,13 @@ function LinearAlgebra.eigen(A::Symmetric{<:Dual{Tg,T,N}, <:StaticArrays.StaticM end # Gradient -@inline ForwardDiff.gradient(f, x::StaticArray) = vector_mode_gradient(f, x) -@inline ForwardDiff.gradient(f, x::StaticArray, cfg::GradientConfig) = gradient(f, x) -@inline ForwardDiff.gradient(f, x::StaticArray, cfg::GradientConfig, ::Val) = gradient(f, x) +@inline ForwardDiff.gradient(f::F, x::StaticArray) where F = vector_mode_gradient(f, x) +@inline ForwardDiff.gradient(f::F, x::StaticArray, cfg::GradientConfig) where F = gradient(f, x) +@inline ForwardDiff.gradient(f::F, x::StaticArray, cfg::GradientConfig, ::Val) where F = gradient(f, x) -@inline ForwardDiff.gradient!(result::Union{AbstractArray,DiffResult}, f, x::StaticArray) = vector_mode_gradient!(result, f, x) -@inline ForwardDiff.gradient!(result::Union{AbstractArray,DiffResult}, f, x::StaticArray, cfg::GradientConfig) = gradient!(result, f, x) -@inline ForwardDiff.gradient!(result::Union{AbstractArray,DiffResult}, f, x::StaticArray, cfg::GradientConfig, ::Val) = gradient!(result, f, x) +@inline ForwardDiff.gradient!(result::Union{AbstractArray,DiffResult}, f::F, x::StaticArray) where F = vector_mode_gradient!(result, f, x) +@inline ForwardDiff.gradient!(result::Union{AbstractArray,DiffResult}, f::F, x::StaticArray, cfg::GradientConfig) where F = gradient!(result, f, x) +@inline ForwardDiff.gradient!(result::Union{AbstractArray,DiffResult}, f::F, x::StaticArray, cfg::GradientConfig, ::Val) where F = gradient!(result, f, x) @generated function extract_gradient(::Type{T}, y::Real, x::S) where {T,S<:StaticArray} result = Expr(:tuple, [:(partials(T, y, $i)) for i in 1:length(x)]...) @@ -65,13 +65,13 @@ end end # Jacobian -@inline ForwardDiff.jacobian(f, x::StaticArray) = vector_mode_jacobian(f, x) -@inline ForwardDiff.jacobian(f, x::StaticArray, cfg::JacobianConfig) = jacobian(f, x) -@inline ForwardDiff.jacobian(f, x::StaticArray, cfg::JacobianConfig, ::Val) = jacobian(f, x) +@inline ForwardDiff.jacobian(f::F, x::StaticArray) where F = vector_mode_jacobian(f, x) +@inline ForwardDiff.jacobian(f::F, x::StaticArray, cfg::JacobianConfig) where F = jacobian(f, x) +@inline ForwardDiff.jacobian(f::F, x::StaticArray, cfg::JacobianConfig, ::Val) where F = jacobian(f, x) -@inline ForwardDiff.jacobian!(result::Union{AbstractArray,DiffResult}, f, x::StaticArray) = vector_mode_jacobian!(result, f, x) -@inline ForwardDiff.jacobian!(result::Union{AbstractArray,DiffResult}, f, x::StaticArray, cfg::JacobianConfig) = jacobian!(result, f, x) -@inline ForwardDiff.jacobian!(result::Union{AbstractArray,DiffResult}, f, x::StaticArray, cfg::JacobianConfig, ::Val) = jacobian!(result, f, x) +@inline ForwardDiff.jacobian!(result::Union{AbstractArray,DiffResult}, f::F, x::StaticArray) where F = vector_mode_jacobian!(result, f, x) +@inline ForwardDiff.jacobian!(result::Union{AbstractArray,DiffResult}, f::F, x::StaticArray, cfg::JacobianConfig) where F = jacobian!(result, f, x) +@inline ForwardDiff.jacobian!(result::Union{AbstractArray,DiffResult}, f::F, x::StaticArray, cfg::JacobianConfig, ::Val) where F = jacobian!(result, f, x) @generated function extract_jacobian(::Type{T}, ydual::StaticArray, x::S) where {T,S<:StaticArray} M, N = length(ydual), length(x) @@ -110,18 +110,18 @@ end end # Hessian -ForwardDiff.hessian(f, x::StaticArray) = jacobian(y -> gradient(f, y), x) -ForwardDiff.hessian(f, x::StaticArray, cfg::HessianConfig) = hessian(f, x) -ForwardDiff.hessian(f, x::StaticArray, cfg::HessianConfig, ::Val) = hessian(f, x) +ForwardDiff.hessian(f::F, x::StaticArray) where F = jacobian(y -> gradient(f, y), x) +ForwardDiff.hessian(f::F, x::StaticArray, cfg::HessianConfig) where F = hessian(f, x) +ForwardDiff.hessian(f::F, x::StaticArray, cfg::HessianConfig, ::Val) where F = hessian(f, x) -ForwardDiff.hessian!(result::AbstractArray, f, x::StaticArray) = jacobian!(result, y -> gradient(f, y), x) +ForwardDiff.hessian!(result::AbstractArray, f::F, x::StaticArray) where F = jacobian!(result, y -> gradient(f, y), x) -ForwardDiff.hessian!(result::MutableDiffResult, f, x::StaticArray) = hessian!(result, f, x, HessianConfig(f, result, x)) +ForwardDiff.hessian!(result::MutableDiffResult, f::F, x::StaticArray) where F = hessian!(result, f, x, HessianConfig(f, result, x)) -ForwardDiff.hessian!(result::ImmutableDiffResult, f, x::StaticArray, cfg::HessianConfig) = hessian!(result, f, x) -ForwardDiff.hessian!(result::ImmutableDiffResult, f, x::StaticArray, cfg::HessianConfig, ::Val) = hessian!(result, f, x) +ForwardDiff.hessian!(result::ImmutableDiffResult, f::F, x::StaticArray, cfg::HessianConfig) where F = hessian!(result, f, x) +ForwardDiff.hessian!(result::ImmutableDiffResult, f::F, x::StaticArray, cfg::HessianConfig, ::Val) where F = hessian!(result, f, x) -function ForwardDiff.hessian!(result::ImmutableDiffResult, f, x::StaticArray) +function ForwardDiff.hessian!(result::ImmutableDiffResult, f::F, x::StaticArray) where F T = typeof(Tag(f, eltype(x))) d1 = dualize(T, x) d2 = dualize(T, d1) diff --git a/src/derivative.jl b/src/derivative.jl index accc2405..b39e2a48 100644 --- a/src/derivative.jl +++ b/src/derivative.jl @@ -22,8 +22,8 @@ stored in `y`. Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care. """ -@inline function derivative(f!, y::AbstractArray, x::Real, - cfg::DerivativeConfig{T} = DerivativeConfig(f!, y, x), ::Val{CHK}=Val{true}()) where {T, CHK} +@inline function derivative(f!::F, y::AbstractArray, x::Real, + cfg::DerivativeConfig{T} = DerivativeConfig(f!, y, x), ::Val{CHK}=Val{true}()) where {F, T, CHK} require_one_based_indexing(y) CHK && checktag(T, f!, x) ydual = cfg.duals @@ -60,8 +60,8 @@ called as `f!(y, x)` where the result is stored in `y`. Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care. """ @inline function derivative!(result::Union{AbstractArray,DiffResult}, - f!, y::AbstractArray, x::Real, - cfg::DerivativeConfig{T} = DerivativeConfig(f!, y, x), ::Val{CHK}=Val{true}()) where {T, CHK} + f!::F, y::AbstractArray, x::Real, + cfg::DerivativeConfig{T} = DerivativeConfig(f!, y, x), ::Val{CHK}=Val{true}()) where {F, T, CHK} result isa DiffResult ? require_one_based_indexing(y) : require_one_based_indexing(result, y) CHK && checktag(T, f!, x) ydual = cfg.duals diff --git a/src/gradient.jl b/src/gradient.jl index c0b4d709..6d4cc791 100644 --- a/src/gradient.jl +++ b/src/gradient.jl @@ -13,7 +13,7 @@ This method assumes that `isa(f(x), Real)`. Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care. """ -function gradient(f, x::AbstractArray, cfg::GradientConfig{T} = GradientConfig(f, x), ::Val{CHK}=Val{true}()) where {T, CHK} +function gradient(f::F, x::AbstractArray, cfg::GradientConfig{T} = GradientConfig(f, x), ::Val{CHK}=Val{true}()) where {F, T, CHK} require_one_based_indexing(x) CHK && checktag(T, f, x) if chunksize(cfg) == length(x) diff --git a/src/hessian.jl b/src/hessian.jl index ababac0f..9c755c9a 100644 --- a/src/hessian.jl +++ b/src/hessian.jl @@ -11,7 +11,7 @@ This method assumes that `isa(f(x), Real)`. Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care. """ -function hessian(f, x::AbstractArray, cfg::HessianConfig{T} = HessianConfig(f, x), ::Val{CHK}=Val{true}()) where {T,CHK} +function hessian(f::F, x::AbstractArray, cfg::HessianConfig{T} = HessianConfig(f, x), ::Val{CHK}=Val{true}()) where {F, T,CHK} require_one_based_indexing(x) CHK && checktag(T, f, x) ∇f = y -> gradient(f, y, cfg.gradient_config, Val{false}()) @@ -28,7 +28,7 @@ This method assumes that `isa(f(x), Real)`. Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care. """ -function hessian!(result::AbstractArray, f, x::AbstractArray, cfg::HessianConfig{T} = HessianConfig(f, x), ::Val{CHK}=Val{true}()) where {T,CHK} +function hessian!(result::AbstractArray, f::F, x::AbstractArray, cfg::HessianConfig{T} = HessianConfig(f, x), ::Val{CHK}=Val{true}()) where {F,T,CHK} require_one_based_indexing(result, x) CHK && checktag(T, f, x) ∇f = y -> gradient(f, y, cfg.gradient_config, Val{false}()) @@ -63,8 +63,7 @@ because `isa(result, DiffResult)`, `cfg` is constructed as `HessianConfig(f, res Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care. """ -function hessian!(result::DiffResult, f, x::AbstractArray, cfg::HessianConfig{T} = HessianConfig(f, result, x), ::Val{CHK}=Val{true}()) where {T,CHK} - require_one_based_indexing(x) +function hessian!(result::DiffResult, f::F, x::AbstractArray, cfg::HessianConfig{T} = HessianConfig(f, result, x), ::Val{CHK}=Val{true}()) where {F,T,CHK} CHK && checktag(T, f, x) ∇f! = InnerGradientForHess(result, cfg, f) jacobian!(DiffResults.hessian(result), ∇f!, DiffResults.gradient(result), x, cfg.jacobian_config, Val{false}()) diff --git a/src/jacobian.jl b/src/jacobian.jl index 1ed3d7b5..e7ca96f5 100644 --- a/src/jacobian.jl +++ b/src/jacobian.jl @@ -15,7 +15,7 @@ This method assumes that `isa(f(x), AbstractArray)`. Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care. """ -function jacobian(f, x::AbstractArray, cfg::JacobianConfig{T} = JacobianConfig(f, x), ::Val{CHK}=Val{true}()) where {T,CHK} +function jacobian(f::F, x::AbstractArray, cfg::JacobianConfig{T} = JacobianConfig(f, x), ::Val{CHK}=Val{true}()) where {F,T,CHK} require_one_based_indexing(x) CHK && checktag(T, f, x) if chunksize(cfg) == length(x) @@ -33,7 +33,7 @@ stored in `y`. Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care. """ -function jacobian(f!, y::AbstractArray, x::AbstractArray, cfg::JacobianConfig{T} = JacobianConfig(f!, y, x), ::Val{CHK}=Val{true}()) where {T, CHK} +function jacobian(f!::F, y::AbstractArray, x::AbstractArray, cfg::JacobianConfig{T} = JacobianConfig(f!, y, x), ::Val{CHK}=Val{true}()) where {F,T, CHK} require_one_based_indexing(y, x) CHK && checktag(T, f!, x) if chunksize(cfg) == length(x) @@ -54,7 +54,7 @@ This method assumes that `isa(f(x), AbstractArray)`. Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care. """ -function jacobian!(result::Union{AbstractArray,DiffResult}, f, x::AbstractArray, cfg::JacobianConfig{T} = JacobianConfig(f, x), ::Val{CHK}=Val{true}()) where {T, CHK} +function jacobian!(result::Union{AbstractArray,DiffResult}, f::F, x::AbstractArray, cfg::JacobianConfig{T} = JacobianConfig(f, x), ::Val{CHK}=Val{true}()) where {F,T, CHK} result isa DiffResult ? require_one_based_indexing(x) : require_one_based_indexing(result, x) CHK && checktag(T, f, x) if chunksize(cfg) == length(x) @@ -75,7 +75,7 @@ This method assumes that `isa(f(x), AbstractArray)`. Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care. """ -function jacobian!(result::Union{AbstractArray,DiffResult}, f!, y::AbstractArray, x::AbstractArray, cfg::JacobianConfig{T} = JacobianConfig(f!, y, x), ::Val{CHK}=Val{true}()) where {T,CHK} +function jacobian!(result::Union{AbstractArray,DiffResult}, f!::F, y::AbstractArray, x::AbstractArray, cfg::JacobianConfig{T} = JacobianConfig(f!, y, x), ::Val{CHK}=Val{true}()) where {F,T,CHK} result isa DiffResult ? require_one_based_indexing(y, x) : require_one_based_indexing(result, y, x) CHK && checktag(T, f!, x) if chunksize(cfg) == length(x)