Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

make StaticArrays dependency into an extension on v0.11-dev #628

Merged
merged 5 commits into from
Feb 25, 2023
Merged

make StaticArrays dependency into an extension on v0.11-dev #628

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
04f76ca
make StaticArrays dependency into an extension
Dec 13, 2022
0632329
update extension naming
ChrisRackauckas Feb 19, 2023
77d9ac2
Make the two extension PRs "exactly" the same
ChrisRackauckas Feb 19, 2023
7cd21b9
Fix extension naming
ChrisRackauckas Feb 19, 2023
df7e6e0
Only import ForwardDiff and StaticArrays
ChrisRackauckas Feb 19, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 8 additions & 1 deletion Project.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,12 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"

[weakdeps]
StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"

[extensions]
ForwardDiffStaticArraysExt = "StaticArrays"

[compat]
Calculus = "0.5"
CommonSubexpressions = "0.3"
Expand All @@ -33,7 +39,8 @@ Calculus = "49dc2e85-a5d0-5ad3-a950-438e2897f1b9"
DiffTests = "de460e47-3fe3-5279-bb4a-814414816d5d"
InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

[targets]
test = ["Calculus", "DiffTests", "SparseArrays", "Test", "InteractiveUtils"]
test = ["Calculus", "DiffTests", "SparseArrays", "StaticArrays", "Test", "InteractiveUtils"]
138 changes: 138 additions & 0 deletions ext/ForwardDiffStaticArraysExt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,138 @@
module ForwardDiffStaticArraysExt

using ForwardDiff, StaticArrays
using ForwardDiff.LinearAlgebra
using ForwardDiff.DiffResults
using ForwardDiff: Dual, partials, GradientConfig, JacobianConfig, HessianConfig, Tag, Chunk,
gradient, hessian, jacobian, gradient!, hessian!, jacobian!,
extract_gradient!, extract_jacobian!, extract_value!,
vector_mode_gradient, vector_mode_gradient!,
vector_mode_jacobian, vector_mode_jacobian!, valtype, value, _lyap_div!
using DiffResults: DiffResult, ImmutableDiffResult, MutableDiffResult

@generated function dualize(::Type{T}, x::StaticArray) where T
N = length(x)
dx = Expr(:tuple, [:(Dual{T}(x[$i], chunk, Val{$i}())) for i in 1:N]...)
V = StaticArrays.similar_type(x, Dual{T,eltype(x),N})
return quote
chunk = Chunk{$N}()
$(Expr(:meta, :inline))
return $V($(dx))
end
end

@inline static_dual_eval(::Type{T}, f, x::StaticArray) where T = f(dualize(T, x))

function LinearAlgebra.eigvals(A::Symmetric{<:Dual{Tg,T,N}, <:StaticArrays.StaticMatrix}) where {Tg,T<:Real,N}
λ,Q = eigen(Symmetric(value.(parent(A))))
parts = ntuple(j -> diag(Q' * getindex.(partials.(A), j) * Q), N)
Dual{Tg}.(λ, tuple.(parts...))
end

function LinearAlgebra.eigen(A::Symmetric{<:Dual{Tg,T,N}, <:StaticArrays.StaticMatrix}) where {Tg,T<:Real,N}
λ = eigvals(A)
_,Q = eigen(Symmetric(value.(parent(A))))
parts = ntuple(j -> Q*_lyap_div!(Q' * getindex.(partials.(A), j) * Q - Diagonal(getindex.(partials.(λ), j)), value.(λ)), N)
Eigen(λ,Dual{Tg}.(Q, tuple.(parts...)))
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!(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)

@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)]...)
return quote
$(Expr(:meta, :inline))
V = StaticArrays.similar_type(S, valtype($y))
return V($result)
end
end

@inline function ForwardDiff.vector_mode_gradient(f, x::StaticArray)
T = typeof(Tag(f, eltype(x)))
return extract_gradient(T, static_dual_eval(T, f, x), x)
end

@inline function ForwardDiff.vector_mode_gradient!(result, f, x::StaticArray)
T = typeof(Tag(f, eltype(x)))
return extract_gradient!(T, result, static_dual_eval(T, f, x))
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!(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)

@generated function extract_jacobian(::Type{T}, ydual::StaticArray, x::S) where {T,S<:StaticArray}
M, N = length(ydual), length(x)
result = Expr(:tuple, [:(partials(T, ydual[$i], $j)) for i in 1:M, j in 1:N]...)
return quote
$(Expr(:meta, :inline))
V = StaticArrays.similar_type(S, valtype(eltype($ydual)), Size($M, $N))
return V($result)
end
end

@inline function ForwardDiff.vector_mode_jacobian(f, x::StaticArray)
T = typeof(Tag(f, eltype(x)))
return extract_jacobian(T, static_dual_eval(T, f, x), x)
end

function extract_jacobian(::Type{T}, ydual::AbstractArray, x::StaticArray) where T
result = similar(ydual, valtype(eltype(ydual)), length(ydual), length(x))
return extract_jacobian!(T, result, ydual, length(x))
end

@inline function ForwardDiff.vector_mode_jacobian!(result, f, x::StaticArray)
T = typeof(Tag(f, eltype(x)))
ydual = static_dual_eval(T, f, x)
result = extract_jacobian!(T, result, ydual, length(x))
result = extract_value!(T, result, ydual)
return result
end

@inline function ForwardDiff.vector_mode_jacobian!(result::ImmutableDiffResult, f, x::StaticArray)
T = typeof(Tag(f, eltype(x)))
ydual = static_dual_eval(T, f, x)
result = DiffResults.jacobian!(result, extract_jacobian(T, ydual, x))
result = DiffResults.value!(d -> value(T,d), result, ydual)
return result
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!(result::AbstractArray, f, x::StaticArray) = 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::ImmutableDiffResult, f, x::StaticArray, cfg::HessianConfig) = hessian!(result, f, x)
ForwardDiff.hessian!(result::ImmutableDiffResult, f, x::StaticArray, cfg::HessianConfig, ::Val) = hessian!(result, f, x)

function ForwardDiff.hessian!(result::ImmutableDiffResult, f, x::StaticArray)
T = typeof(Tag(f, eltype(x)))
d1 = dualize(T, x)
d2 = dualize(T, d1)
fd2 = f(d2)
val = value(T,value(T,fd2))
grad = extract_gradient(T,value(T,fd2), x)
hess = extract_jacobian(T,partials(T,fd2), x)
result = DiffResults.hessian!(result, hess)
result = DiffResults.gradient!(result, grad)
result = DiffResults.value!(result, val)
return result
end

end
7 changes: 5 additions & 2 deletions src/ForwardDiff.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,7 @@
module ForwardDiff

using DiffRules, DiffResults
using DiffResults: DiffResult, MutableDiffResult, ImmutableDiffResult
using StaticArrays
using DiffResults: DiffResult, MutableDiffResult
using Preferences
using Random
using LinearAlgebra
Expand All @@ -23,6 +22,10 @@ include("gradient.jl")
include("jacobian.jl")
include("hessian.jl")

if !isdefined(Base, :get_extension)
include("../ext/ForwardDiffStaticArraysExt.jl")
end

export DiffResults

end # module
13 changes: 0 additions & 13 deletions src/apiutils.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,19 +18,6 @@ end
# vector mode function evaluation #
###################################

@generated function dualize(::Type{T}, x::StaticArray) where T
N = length(x)
dx = Expr(:tuple, [:(Dual{T}(x[$i], chunk, Val{$i}())) for i in 1:N]...)
V = StaticArrays.similar_type(x, Dual{T,eltype(x),N})
return quote
chunk = Chunk{$N}()
$(Expr(:meta, :inline))
return $V($(dx))
end
end

@inline static_dual_eval(::Type{T}, f, x::StaticArray) where T = f(dualize(T, x))

function vector_mode_dual_eval!(f::F, cfg::Union{JacobianConfig,GradientConfig}, x) where {F}
xdual = cfg.duals
seed!(xdual, x, cfg.seeds)
Expand Down
13 changes: 0 additions & 13 deletions src/dual.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -726,12 +726,6 @@ function LinearAlgebra.eigvals(A::Symmetric{<:Dual{Tg,T,N}}) where {Tg,T<:Real,N
Dual{Tg}.(λ, tuple.(parts...))
end

function LinearAlgebra.eigvals(A::Symmetric{<:Dual{Tg,T,N}, <:StaticArrays.StaticMatrix}) where {Tg,T<:Real,N}
λ,Q = eigen(Symmetric(value.(parent(A))))
parts = ntuple(j -> diag(Q' * getindex.(partials.(A), j) * Q), N)
Dual{Tg}.(λ, tuple.(parts...))
end

function LinearAlgebra.eigvals(A::Hermitian{<:Complex{<:Dual{Tg,T,N}}}) where {Tg,T<:Real,N}
λ,Q = eigen(Hermitian(value.(real.(parent(A))) .+ im .* value.(imag.(parent(A)))))
parts = ntuple(j -> diag(real.(Q' * (getindex.(partials.(real.(A)) .+ im .* partials.(imag.(A)), j)) * Q)), N)
Expand Down Expand Up @@ -761,13 +755,6 @@ function LinearAlgebra.eigen(A::Symmetric{<:Dual{Tg,T,N}}) where {Tg,T<:Real,N}
Eigen(λ,Dual{Tg}.(Q, tuple.(parts...)))
end

function LinearAlgebra.eigen(A::Symmetric{<:Dual{Tg,T,N}, <:StaticArrays.StaticMatrix}) where {Tg,T<:Real,N}
λ = eigvals(A)
_,Q = eigen(Symmetric(value.(parent(A))))
parts = ntuple(j -> Q*_lyap_div!(Q' * getindex.(partials.(A), j) * Q - Diagonal(getindex.(partials.(λ), j)), value.(λ)), N)
Eigen(λ,Dual{Tg}.(Q, tuple.(parts...)))
end

function LinearAlgebra.eigen(A::SymTridiagonal{<:Dual{Tg,T,N}}) where {Tg,T<:Real,N}
λ = eigvals(A)
_,Q = eigen(SymTridiagonal(value.(parent(A))))
Expand Down
27 changes: 0 additions & 27 deletions src/gradient.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,29 +43,12 @@ function gradient!(result::Union{AbstractArray,DiffResult}, f::F, x::AbstractArr
return result
end

@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::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)?"))

#####################
# result extraction #
#####################

@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)]...)
return quote
$(Expr(:meta, :inline))
V = StaticArrays.similar_type(S, valtype($y))
return V($result)
end
end

function extract_gradient!(::Type{T}, result::DiffResult, y::Real) where {T}
result = DiffResults.value!(result, y)
grad = DiffResults.gradient(result)
Expand Down Expand Up @@ -117,16 +100,6 @@ function vector_mode_gradient!(result, f::F, x, cfg::GradientConfig{T}) where {T
return result
end

@inline function vector_mode_gradient(f, x::StaticArray)
T = typeof(Tag(f, eltype(x)))
return extract_gradient(T, static_dual_eval(T, f, x), x)
end

@inline function vector_mode_gradient!(result, f, x::StaticArray)
T = typeof(Tag(f, eltype(x)))
return extract_gradient!(T, result, static_dual_eval(T, f, x))
end

##############
# chunk mode #
##############
Expand Down
25 changes: 0 additions & 25 deletions src/hessian.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,28 +69,3 @@ function hessian!(result::DiffResult, f::F, x::AbstractArray, cfg::HessianConfig
jacobian!(DiffResults.hessian(result), ∇f!, DiffResults.gradient(result), x, cfg.jacobian_config, Val{false}())
return ∇f!.result
end

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::F, x::StaticArray) where F = jacobian!(result, y -> gradient(f, y), x)

hessian!(result::MutableDiffResult, f::F, x::StaticArray) where F = hessian!(result, f, x, HessianConfig(f, result, 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::F, x::StaticArray) where F
T = typeof(Tag(f, eltype(x)))
d1 = dualize(T, x)
d2 = dualize(T, d1)
fd2 = f(d2)
val = value(T,value(T,fd2))
grad = extract_gradient(T,value(T,fd2), x)
hess = extract_jacobian(T,partials(T,fd2), x)
result = DiffResults.hessian!(result, hess)
result = DiffResults.gradient!(result, grad)
result = DiffResults.value!(result, val)
return result
end
44 changes: 0 additions & 44 deletions src/jacobian.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -86,35 +86,12 @@ function jacobian!(result::Union{AbstractArray,DiffResult}, f!::F, y::AbstractAr
return result
end

@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::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)?"))

#####################
# result extraction #
#####################

@generated function extract_jacobian(::Type{T}, ydual::StaticArray, x::S) where {T,S<:StaticArray}
M, N = length(ydual), length(x)
result = Expr(:tuple, [:(partials(T, ydual[$i], $j)) for i in 1:M, j in 1:N]...)
return quote
$(Expr(:meta, :inline))
V = StaticArrays.similar_type(S, valtype(eltype($ydual)), Size($M, $N))
return V($result)
end
end

function extract_jacobian(::Type{T}, ydual::AbstractArray, x::StaticArray) where T
result = similar(ydual, valtype(eltype(ydual)), length(ydual), length(x))
return extract_jacobian!(T, result, ydual, length(x))
end

function extract_jacobian!(::Type{T}, result::AbstractArray, ydual::AbstractArray, n) where {T}
out_reshaped = reshape(result, length(ydual), n)
ydual_reshaped = vec(ydual)
Expand Down Expand Up @@ -184,27 +161,6 @@ function vector_mode_jacobian!(result, f!::F, y, x, cfg::JacobianConfig{T}) wher
return result
end

@inline function vector_mode_jacobian(f, x::StaticArray)
T = typeof(Tag(f, eltype(x)))
return extract_jacobian(T, static_dual_eval(T, f, x), x)
end

@inline function vector_mode_jacobian!(result, f, x::StaticArray)
T = typeof(Tag(f, eltype(x)))
ydual = static_dual_eval(T, f, x)
result = extract_jacobian!(T, result, ydual, length(x))
result = extract_value!(T, result, ydual)
return result
end

@inline function vector_mode_jacobian!(result::ImmutableDiffResult, f, x::StaticArray)
T = typeof(Tag(f, eltype(x)))
ydual = static_dual_eval(T, f, x)
result = DiffResults.jacobian!(result, extract_jacobian(T, ydual, x))
result = DiffResults.value!(d -> value(T,d), result, ydual)
return result
end

const JACOBIAN_ERROR = DimensionMismatch("jacobian(f, x) expects that f(x) is an array. Perhaps you meant gradient(f, x)?")

# chunk mode #
Expand Down