From 098311b9ea5a292371d29d9d864b165ce60e0f76 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Fuhrmann?= Date: Thu, 8 Dec 2022 18:51:25 +0100 Subject: [PATCH 1/2] specialize on function in jacobian --- src/jacobian.jl | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/src/jacobian.jl b/src/jacobian.jl index bcda61d7..a1a05c10 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} CHK && checktag(T, f, x) if chunksize(cfg) == length(x) return vector_mode_jacobian(f, x, cfg) @@ -32,7 +32,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} CHK && checktag(T, f!, x) if chunksize(cfg) == length(x) return vector_mode_jacobian(f!, y, x, cfg) @@ -52,7 +52,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} CHK && checktag(T, f, x) if chunksize(cfg) == length(x) vector_mode_jacobian!(result, f, x, cfg) @@ -72,7 +72,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} CHK && checktag(T, f!, x) if chunksize(cfg) == length(x) vector_mode_jacobian!(result, f!, y, x, cfg) @@ -82,13 +82,13 @@ function jacobian!(result::Union{AbstractArray,DiffResult}, f!, y::AbstractArray return result end -@inline jacobian(f, x::StaticArray) = vector_mode_jacobian(f, x) -@inline jacobian(f, x::StaticArray, cfg::JacobianConfig) = jacobian(f, x) -@inline jacobian(f, x::StaticArray, cfg::JacobianConfig, ::Val) = jacobian(f, x) +@inline jacobian(f::F, x::StaticArray) where F = vector_mode_jacobian(f, x) +@inline jacobian(f::F, x::StaticArray, cfg::JacobianConfig) where F = jacobian(f, x) +@inline jacobian(f::F, x::StaticArray, cfg::JacobianConfig, ::Val) where F = jacobian(f, x) -@inline jacobian!(result::Union{AbstractArray,DiffResult}, f, x::StaticArray) = vector_mode_jacobian!(result, f, x) -@inline jacobian!(result::Union{AbstractArray,DiffResult}, f, x::StaticArray, cfg::JacobianConfig) = jacobian!(result, f, x) -@inline jacobian!(result::Union{AbstractArray,DiffResult}, f, x::StaticArray, cfg::JacobianConfig, ::Val) = jacobian!(result, f, x) +@inline jacobian!(result::Union{AbstractArray,DiffResult}, f::F, x::StaticArray) where F = vector_mode_jacobian!(result, f, x) +@inline jacobian!(result::Union{AbstractArray,DiffResult}, f::F, x::StaticArray, cfg::JacobianConfig) where F = jacobian!(result, f, x) +@inline jacobian!(result::Union{AbstractArray,DiffResult}, f::F, x::StaticArray, cfg::JacobianConfig, ::Val) where F = jacobian!(result, f, x) jacobian(f, x::Real) = throw(DimensionMismatch("jacobian(f, x) expects that x is an array. Perhaps you meant derivative(f, x)?")) From 4b310fdc89729690295a18afa1ce75ede8ba809d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Fuhrmann?= Date: Thu, 8 Dec 2022 19:03:08 +0100 Subject: [PATCH 2/2] specilize on function parameters for derivative, gradient, hessian --- src/derivative.jl | 8 ++++---- src/gradient.jl | 14 +++++++------- src/hessian.jl | 22 +++++++++++----------- 3 files changed, 22 insertions(+), 22 deletions(-) diff --git a/src/derivative.jl b/src/derivative.jl index 5d7e5dab..9ca22f3c 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} CHK && checktag(T, f!, x) ydual = cfg.duals seed!(ydual, y) @@ -58,8 +58,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} CHK && checktag(T, f!, x) ydual = cfg.duals seed!(ydual, y) diff --git a/src/gradient.jl b/src/gradient.jl index f9f173eb..f45869d9 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} CHK && checktag(T, f, x) if chunksize(cfg) == length(x) return vector_mode_gradient(f, x, cfg) @@ -41,13 +41,13 @@ function gradient!(result::Union{AbstractArray,DiffResult}, f::F, x::AbstractArr return result end -@inline gradient(f, x::StaticArray) = vector_mode_gradient(f, x) -@inline gradient(f, x::StaticArray, cfg::GradientConfig) = gradient(f, x) -@inline gradient(f, x::StaticArray, cfg::GradientConfig, ::Val) = gradient(f, x) +@inline gradient(f::F, x::StaticArray) where F = vector_mode_gradient(f, x) +@inline gradient(f::F, x::StaticArray, cfg::GradientConfig) where F = gradient(f, x) +@inline gradient(f::F, x::StaticArray, cfg::GradientConfig, ::Val) where F = gradient(f, x) -@inline gradient!(result::Union{AbstractArray,DiffResult}, f, x::StaticArray) = vector_mode_gradient!(result, f, x) -@inline gradient!(result::Union{AbstractArray,DiffResult}, f, x::StaticArray, cfg::GradientConfig) = gradient!(result, f, x) -@inline gradient!(result::Union{AbstractArray,DiffResult}, f, x::StaticArray, cfg::GradientConfig, ::Val) = gradient!(result, f, x) +@inline gradient!(result::Union{AbstractArray,DiffResult}, f::F, x::StaticArray) where F = vector_mode_gradient!(result, f, x) +@inline gradient!(result::Union{AbstractArray,DiffResult}, f::F, x::StaticArray, cfg::GradientConfig) where F = gradient!(result, f, x) +@inline gradient!(result::Union{AbstractArray,DiffResult}, f::F, x::StaticArray, cfg::GradientConfig, ::Val) where F = gradient!(result, f, x) gradient(f, x::Real) = throw(DimensionMismatch("gradient(f, x) expects that x is an array. Perhaps you meant derivative(f, x)?")) diff --git a/src/hessian.jl b/src/hessian.jl index 7de6109d..8fcc793e 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} CHK && checktag(T, f, x) ∇f = y -> gradient(f, y, cfg.gradient_config, Val{false}()) return jacobian(∇f, x, cfg.jacobian_config, Val{false}()) @@ -27,7 +27,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} CHK && checktag(T, f, x) ∇f = y -> gradient(f, y, cfg.gradient_config, Val{false}()) jacobian!(result, ∇f, x, cfg.jacobian_config, Val{false}()) @@ -61,25 +61,25 @@ 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} +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}()) return ∇f!.result end -hessian(f, x::StaticArray) = jacobian(y -> gradient(f, y), x) -hessian(f, x::StaticArray, cfg::HessianConfig) = hessian(f, x) -hessian(f, x::StaticArray, cfg::HessianConfig, ::Val) = hessian(f, x) +hessian(f::F, x::StaticArray) where F = jacobian(y -> gradient(f, y), x) +hessian(f::F, x::StaticArray, cfg::HessianConfig) where F = hessian(f, x) +hessian(f::F, x::StaticArray, cfg::HessianConfig, ::Val) where F = hessian(f, x) -hessian!(result::AbstractArray, f, x::StaticArray) = jacobian!(result, y -> gradient(f, y), x) +hessian!(result::AbstractArray, f::F, x::StaticArray) where F = jacobian!(result, y -> gradient(f, y), x) -hessian!(result::MutableDiffResult, f, x::StaticArray) = hessian!(result, f, x, HessianConfig(f, result, x)) +hessian!(result::MutableDiffResult, f::F, x::StaticArray) where F = hessian!(result, f, x, HessianConfig(f, result, x)) -hessian!(result::ImmutableDiffResult, f, x::StaticArray, cfg::HessianConfig) = hessian!(result, f, x) -hessian!(result::ImmutableDiffResult, f, x::StaticArray, cfg::HessianConfig, ::Val) = hessian!(result, f, x) +hessian!(result::ImmutableDiffResult, f::F, x::StaticArray, cfg::HessianConfig) where F = hessian!(result, f, x) +hessian!(result::ImmutableDiffResult, f::F, x::StaticArray, cfg::HessianConfig, ::Val) where F = hessian!(result, f, x) -function hessian!(result::ImmutableDiffResult, f, x::StaticArray) +function hessian!(result::ImmutableDiffResult, f::F, x::StaticArray) where F T = typeof(Tag(f, eltype(x))) d1 = dualize(T, x) d2 = dualize(T, d1)