Skip to content

Commit 136cdbc

Browse files
ChrisRackauckasChrisRackauckas
authored
Merge pull request #625 from JuliaDiff/staticarraysext
Turn StaticArrays into an Extension Package
2 parents eb8d755 + bc3a395 commit 136cdbc

File tree

8 files changed

+151
-125
lines changed

8 files changed

+151
-125
lines changed

‎Project.toml

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,12 +28,19 @@ SpecialFunctions = "0.8, 0.9, 0.10, 1.0, 2"
2828
StaticArrays = "0.8.3, 0.9, 0.10, 0.11, 0.12, 1.0"
2929
julia = "1"
3030

31+
[extensions]
32+
ForwardDiffStaticArraysExt = "StaticArrays"
33+
3134
[extras]
3235
Calculus = "49dc2e85-a5d0-5ad3-a950-438e2897f1b9"
3336
DiffTests = "de460e47-3fe3-5279-bb4a-814414816d5d"
3437
InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
3538
SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
39+
StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
3640
Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
3741

3842
[targets]
39-
test = ["Calculus", "DiffTests", "SparseArrays", "Test", "InteractiveUtils"]
43+
test = ["Calculus", "DiffTests", "SparseArrays", "Test", "InteractiveUtils", "StaticArrays"]
44+
45+
[weakdeps]
46+
StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"

‎ext/ForwardDiffStaticArraysExt.jl

Lines changed: 138 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,138 @@
1+
module ForwardDiffStaticArraysExt
2+
3+
using ForwardDiff, StaticArrays
4+
using ForwardDiff.LinearAlgebra
5+
using ForwardDiff.DiffResults
6+
using ForwardDiff: Dual, partials, GradientConfig, JacobianConfig, HessianConfig, Tag, Chunk,
7+
gradient, hessian, jacobian, gradient!, hessian!, jacobian!,
8+
extract_gradient!, extract_jacobian!, extract_value!,
9+
vector_mode_gradient, vector_mode_gradient!,
10+
vector_mode_jacobian, vector_mode_jacobian!, valtype, value, _lyap_div!
11+
using DiffResults: DiffResult, ImmutableDiffResult, MutableDiffResult
12+
13+
@generated function dualize(::Type{T}, x::StaticArray) where T
14+
N = length(x)
15+
dx = Expr(:tuple, [:(Dual{T}(x[$i], chunk, Val{$i}())) for i in 1:N]...)
16+
V = StaticArrays.similar_type(x, Dual{T,eltype(x),N})
17+
return quote
18+
chunk = Chunk{$N}()
19+
$(Expr(:meta, :inline))
20+
return $V($(dx))
21+
end
22+
end
23+
24+
@inline static_dual_eval(::Type{T}, f, x::StaticArray) where T = f(dualize(T, x))
25+
26+
function LinearAlgebra.eigvals(A::Symmetric{<:Dual{Tg,T,N}, <:StaticArrays.StaticMatrix}) where {Tg,T<:Real,N}
27+
λ,Q = eigen(Symmetric(value.(parent(A))))
28+
parts = ntuple(j -> diag(Q' * getindex.(partials.(A), j) * Q), N)
29+
Dual{Tg}.(λ, tuple.(parts...))
30+
end
31+
32+
function LinearAlgebra.eigen(A::Symmetric{<:Dual{Tg,T,N}, <:StaticArrays.StaticMatrix}) where {Tg,T<:Real,N}
33+
λ = eigvals(A)
34+
_,Q = eigen(Symmetric(value.(parent(A))))
35+
parts = ntuple(j -> Q*ForwardDiff._lyap_div!(Q' * getindex.(partials.(A), j) * Q - Diagonal(getindex.(partials.(λ), j)), value.(λ)), N)
36+
Eigen(λ,Dual{Tg}.(Q, tuple.(parts...)))
37+
end
38+
39+
# Gradient
40+
@inline ForwardDiff.gradient(f, x::StaticArray) = vector_mode_gradient(f, x)
41+
@inline ForwardDiff.gradient(f, x::StaticArray, cfg::GradientConfig) = gradient(f, x)
42+
@inline ForwardDiff.gradient(f, x::StaticArray, cfg::GradientConfig, ::Val) = gradient(f, x)
43+
44+
@inline ForwardDiff.gradient!(result::Union{AbstractArray,DiffResult}, f, x::StaticArray) = vector_mode_gradient!(result, f, x)
45+
@inline ForwardDiff.gradient!(result::Union{AbstractArray,DiffResult}, f, x::StaticArray, cfg::GradientConfig) = gradient!(result, f, x)
46+
@inline ForwardDiff.gradient!(result::Union{AbstractArray,DiffResult}, f, x::StaticArray, cfg::GradientConfig, ::Val) = gradient!(result, f, x)
47+
48+
@generated function extract_gradient(::Type{T}, y::Real, x::S) where {T,S<:StaticArray}
49+
result = Expr(:tuple, [:(partials(T, y, $i)) for i in 1:length(x)]...)
50+
return quote
51+
$(Expr(:meta, :inline))
52+
V = StaticArrays.similar_type(S, valtype($y))
53+
return V($result)
54+
end
55+
end
56+
57+
@inline function ForwardDiff.vector_mode_gradient(f, x::StaticArray)
58+
T = typeof(Tag(f, eltype(x)))
59+
return extract_gradient(T, static_dual_eval(T, f, x), x)
60+
end
61+
62+
@inline function ForwardDiff.vector_mode_gradient!(result, f, x::StaticArray)
63+
T = typeof(Tag(f, eltype(x)))
64+
return extract_gradient!(T, result, static_dual_eval(T, f, x))
65+
end
66+
67+
# Jacobian
68+
@inline ForwardDiff.jacobian(f, x::StaticArray) = vector_mode_jacobian(f, x)
69+
@inline ForwardDiff.jacobian(f, x::StaticArray, cfg::JacobianConfig) = jacobian(f, x)
70+
@inline ForwardDiff.jacobian(f, x::StaticArray, cfg::JacobianConfig, ::Val) = jacobian(f, x)
71+
72+
@inline ForwardDiff.jacobian!(result::Union{AbstractArray,DiffResult}, f, x::StaticArray) = vector_mode_jacobian!(result, f, x)
73+
@inline ForwardDiff.jacobian!(result::Union{AbstractArray,DiffResult}, f, x::StaticArray, cfg::JacobianConfig) = jacobian!(result, f, x)
74+
@inline ForwardDiff.jacobian!(result::Union{AbstractArray,DiffResult}, f, x::StaticArray, cfg::JacobianConfig, ::Val) = jacobian!(result, f, x)
75+
76+
@generated function extract_jacobian(::Type{T}, ydual::StaticArray, x::S) where {T,S<:StaticArray}
77+
M, N = length(ydual), length(x)
78+
result = Expr(:tuple, [:(partials(T, ydual[$i], $j)) for i in 1:M, j in 1:N]...)
79+
return quote
80+
$(Expr(:meta, :inline))
81+
V = StaticArrays.similar_type(S, valtype(eltype($ydual)), Size($M, $N))
82+
return V($result)
83+
end
84+
end
85+
86+
function extract_jacobian(::Type{T}, ydual::AbstractArray, x::StaticArray) where T
87+
result = similar(ydual, valtype(eltype(ydual)), length(ydual), length(x))
88+
return extract_jacobian!(T, result, ydual, length(x))
89+
end
90+
91+
@inline function ForwardDiff.vector_mode_jacobian(f, x::StaticArray)
92+
T = typeof(Tag(f, eltype(x)))
93+
return extract_jacobian(T, static_dual_eval(T, f, x), x)
94+
end
95+
96+
@inline function ForwardDiff.vector_mode_jacobian!(result, f, x::StaticArray)
97+
T = typeof(Tag(f, eltype(x)))
98+
ydual = static_dual_eval(T, f, x)
99+
result = extract_jacobian!(T, result, ydual, length(x))
100+
result = extract_value!(T, result, ydual)
101+
return result
102+
end
103+
104+
@inline function ForwardDiff.vector_mode_jacobian!(result::ImmutableDiffResult, f, x::StaticArray)
105+
T = typeof(Tag(f, eltype(x)))
106+
ydual = static_dual_eval(T, f, x)
107+
result = DiffResults.jacobian!(result, extract_jacobian(T, ydual, x))
108+
result = DiffResults.value!(d -> value(T,d), result, ydual)
109+
return result
110+
end
111+
112+
# Hessian
113+
ForwardDiff.hessian(f, x::StaticArray) = jacobian(y -> gradient(f, y), x)
114+
ForwardDiff.hessian(f, x::StaticArray, cfg::HessianConfig) = hessian(f, x)
115+
ForwardDiff.hessian(f, x::StaticArray, cfg::HessianConfig, ::Val) = hessian(f, x)
116+
117+
ForwardDiff.hessian!(result::AbstractArray, f, x::StaticArray) = jacobian!(result, y -> gradient(f, y), x)
118+
119+
ForwardDiff.hessian!(result::MutableDiffResult, f, x::StaticArray) = hessian!(result, f, x, HessianConfig(f, result, x))
120+
121+
ForwardDiff.hessian!(result::ImmutableDiffResult, f, x::StaticArray, cfg::HessianConfig) = hessian!(result, f, x)
122+
ForwardDiff.hessian!(result::ImmutableDiffResult, f, x::StaticArray, cfg::HessianConfig, ::Val) = hessian!(result, f, x)
123+
124+
function ForwardDiff.hessian!(result::ImmutableDiffResult, f, x::StaticArray)
125+
T = typeof(Tag(f, eltype(x)))
126+
d1 = dualize(T, x)
127+
d2 = dualize(T, d1)
128+
fd2 = f(d2)
129+
val = value(T,value(T,fd2))
130+
grad = extract_gradient(T,value(T,fd2), x)
131+
hess = extract_jacobian(T,partials(T,fd2), x)
132+
result = DiffResults.hessian!(result, hess)
133+
result = DiffResults.gradient!(result, grad)
134+
result = DiffResults.value!(result, val)
135+
return result
136+
end
137+
138+
end

‎src/ForwardDiff.jl

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,7 @@
11
module ForwardDiff
22

33
using DiffRules, DiffResults
4-
using DiffResults: DiffResult, MutableDiffResult, ImmutableDiffResult
5-
using StaticArrays
4+
using DiffResults: DiffResult, MutableDiffResult
65
if VERSION >= v"1.6"
76
using Preferences
87
end
@@ -25,6 +24,10 @@ include("gradient.jl")
2524
include("jacobian.jl")
2625
include("hessian.jl")
2726

27+
if !isdefined(Base, :get_extension)
28+
include("../ext/ForwardDiffStaticArraysExt.jl")
29+
end
30+
2831
export DiffResults
2932

3033
end # module

‎src/apiutils.jl

Lines changed: 0 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -18,19 +18,6 @@ end
1818
# vector mode function evaluation #
1919
###################################
2020

21-
@generated function dualize(::Type{T}, x::StaticArray) where T
22-
N = length(x)
23-
dx = Expr(:tuple, [:(Dual{T}(x[$i], chunk, Val{$i}())) for i in 1:N]...)
24-
V = StaticArrays.similar_type(x, Dual{T,eltype(x),N})
25-
return quote
26-
chunk = Chunk{$N}()
27-
$(Expr(:meta, :inline))
28-
return $V($(dx))
29-
end
30-
end
31-
32-
@inline static_dual_eval(::Type{T}, f, x::StaticArray) where T = f(dualize(T, x))
33-
3421
function vector_mode_dual_eval!(f::F, cfg::Union{JacobianConfig,GradientConfig}, x) where {F}
3522
xdual = cfg.duals
3623
seed!(xdual, x, cfg.seeds)

‎src/dual.jl

Lines changed: 0 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -708,12 +708,6 @@ function LinearAlgebra.eigvals(A::Symmetric{<:Dual{Tg,T,N}}) where {Tg,T<:Real,N
708708
Dual{Tg}.(λ, tuple.(parts...))
709709
end
710710

711-
function LinearAlgebra.eigvals(A::Symmetric{<:Dual{Tg,T,N}, <:StaticArrays.StaticMatrix}) where {Tg,T<:Real,N}
712-
λ,Q = eigen(Symmetric(value.(parent(A))))
713-
parts = ntuple(j -> diag(Q' * getindex.(partials.(A), j) * Q), N)
714-
Dual{Tg}.(λ, tuple.(parts...))
715-
end
716-
717711
function LinearAlgebra.eigvals(A::Hermitian{<:Complex{<:Dual{Tg,T,N}}}) where {Tg,T<:Real,N}
718712
λ,Q = eigen(Hermitian(value.(real.(parent(A))) .+ im .* value.(imag.(parent(A)))))
719713
parts = ntuple(j -> diag(real.(Q' * (getindex.(partials.(real.(A)) .+ im .* partials.(imag.(A)), j)) * Q)), N)
@@ -743,13 +737,6 @@ function LinearAlgebra.eigen(A::Symmetric{<:Dual{Tg,T,N}}) where {Tg,T<:Real,N}
743737
Eigen(λ,Dual{Tg}.(Q, tuple.(parts...)))
744738
end
745739

746-
function LinearAlgebra.eigen(A::Symmetric{<:Dual{Tg,T,N}, <:StaticArrays.StaticMatrix}) where {Tg,T<:Real,N}
747-
λ = eigvals(A)
748-
_,Q = eigen(Symmetric(value.(parent(A))))
749-
parts = ntuple(j -> Q*_lyap_div!(Q' * getindex.(partials.(A), j) * Q - Diagonal(getindex.(partials.(λ), j)), value.(λ)), N)
750-
Eigen(λ,Dual{Tg}.(Q, tuple.(parts...)))
751-
end
752-
753740
function LinearAlgebra.eigen(A::SymTridiagonal{<:Dual{Tg,T,N}}) where {Tg,T<:Real,N}
754741
λ = eigvals(A)
755742
_,Q = eigen(SymTridiagonal(value.(parent(A))))

‎src/gradient.jl

Lines changed: 0 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -41,29 +41,12 @@ function gradient!(result::Union{AbstractArray,DiffResult}, f::F, x::AbstractArr
4141
return result
4242
end
4343

44-
@inline gradient(f, x::StaticArray) = vector_mode_gradient(f, x)
45-
@inline gradient(f, x::StaticArray, cfg::GradientConfig) = gradient(f, x)
46-
@inline gradient(f, x::StaticArray, cfg::GradientConfig, ::Val) = gradient(f, x)
47-
48-
@inline gradient!(result::Union{AbstractArray,DiffResult}, f, x::StaticArray) = vector_mode_gradient!(result, f, x)
49-
@inline gradient!(result::Union{AbstractArray,DiffResult}, f, x::StaticArray, cfg::GradientConfig) = gradient!(result, f, x)
50-
@inline gradient!(result::Union{AbstractArray,DiffResult}, f, x::StaticArray, cfg::GradientConfig, ::Val) = gradient!(result, f, x)
51-
5244
gradient(f, x::Real) = throw(DimensionMismatch("gradient(f, x) expects that x is an array. Perhaps you meant derivative(f, x)?"))
5345

5446
#####################
5547
# result extraction #
5648
#####################
5749

58-
@generated function extract_gradient(::Type{T}, y::Real, x::S) where {T,S<:StaticArray}
59-
result = Expr(:tuple, [:(partials(T, y, $i)) for i in 1:length(x)]...)
60-
return quote
61-
$(Expr(:meta, :inline))
62-
V = StaticArrays.similar_type(S, valtype($y))
63-
return V($result)
64-
end
65-
end
66-
6750
function extract_gradient!(::Type{T}, result::DiffResult, y::Real) where {T}
6851
result = DiffResults.value!(result, y)
6952
grad = DiffResults.gradient(result)
@@ -115,16 +98,6 @@ function vector_mode_gradient!(result, f::F, x, cfg::GradientConfig{T}) where {T
11598
return result
11699
end
117100

118-
@inline function vector_mode_gradient(f, x::StaticArray)
119-
T = typeof(Tag(f, eltype(x)))
120-
return extract_gradient(T, static_dual_eval(T, f, x), x)
121-
end
122-
123-
@inline function vector_mode_gradient!(result, f, x::StaticArray)
124-
T = typeof(Tag(f, eltype(x)))
125-
return extract_gradient!(T, result, static_dual_eval(T, f, x))
126-
end
127-
128101
##############
129102
# chunk mode #
130103
##############

‎src/hessian.jl

Lines changed: 0 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -67,28 +67,3 @@ function hessian!(result::DiffResult, f, x::AbstractArray, cfg::HessianConfig{T}
6767
jacobian!(DiffResults.hessian(result), ∇f!, DiffResults.gradient(result), x, cfg.jacobian_config, Val{false}())
6868
return ∇f!.result
6969
end
70-
71-
hessian(f, x::StaticArray) = jacobian(y -> gradient(f, y), x)
72-
hessian(f, x::StaticArray, cfg::HessianConfig) = hessian(f, x)
73-
hessian(f, x::StaticArray, cfg::HessianConfig, ::Val) = hessian(f, x)
74-
75-
hessian!(result::AbstractArray, f, x::StaticArray) = jacobian!(result, y -> gradient(f, y), x)
76-
77-
hessian!(result::MutableDiffResult, f, x::StaticArray) = hessian!(result, f, x, HessianConfig(f, result, x))
78-
79-
hessian!(result::ImmutableDiffResult, f, x::StaticArray, cfg::HessianConfig) = hessian!(result, f, x)
80-
hessian!(result::ImmutableDiffResult, f, x::StaticArray, cfg::HessianConfig, ::Val) = hessian!(result, f, x)
81-
82-
function hessian!(result::ImmutableDiffResult, f, x::StaticArray)
83-
T = typeof(Tag(f, eltype(x)))
84-
d1 = dualize(T, x)
85-
d2 = dualize(T, d1)
86-
fd2 = f(d2)
87-
val = value(T,value(T,fd2))
88-
grad = extract_gradient(T,value(T,fd2), x)
89-
hess = extract_jacobian(T,partials(T,fd2), x)
90-
result = DiffResults.hessian!(result, hess)
91-
result = DiffResults.gradient!(result, grad)
92-
result = DiffResults.value!(result, val)
93-
return result
94-
end

‎src/jacobian.jl

Lines changed: 0 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -82,35 +82,12 @@ function jacobian!(result::Union{AbstractArray,DiffResult}, f!, y::AbstractArray
8282
return result
8383
end
8484

85-
@inline jacobian(f, x::StaticArray) = vector_mode_jacobian(f, x)
86-
@inline jacobian(f, x::StaticArray, cfg::JacobianConfig) = jacobian(f, x)
87-
@inline jacobian(f, x::StaticArray, cfg::JacobianConfig, ::Val) = jacobian(f, x)
88-
89-
@inline jacobian!(result::Union{AbstractArray,DiffResult}, f, x::StaticArray) = vector_mode_jacobian!(result, f, x)
90-
@inline jacobian!(result::Union{AbstractArray,DiffResult}, f, x::StaticArray, cfg::JacobianConfig) = jacobian!(result, f, x)
91-
@inline jacobian!(result::Union{AbstractArray,DiffResult}, f, x::StaticArray, cfg::JacobianConfig, ::Val) = jacobian!(result, f, x)
92-
9385
jacobian(f, x::Real) = throw(DimensionMismatch("jacobian(f, x) expects that x is an array. Perhaps you meant derivative(f, x)?"))
9486

9587
#####################
9688
# result extraction #
9789
#####################
9890

99-
@generated function extract_jacobian(::Type{T}, ydual::StaticArray, x::S) where {T,S<:StaticArray}
100-
M, N = length(ydual), length(x)
101-
result = Expr(:tuple, [:(partials(T, ydual[$i], $j)) for i in 1:M, j in 1:N]...)
102-
return quote
103-
$(Expr(:meta, :inline))
104-
V = StaticArrays.similar_type(S, valtype(eltype($ydual)), Size($M, $N))
105-
return V($result)
106-
end
107-
end
108-
109-
function extract_jacobian(::Type{T}, ydual::AbstractArray, x::StaticArray) where T
110-
result = similar(ydual, valtype(eltype(ydual)), length(ydual), length(x))
111-
return extract_jacobian!(T, result, ydual, length(x))
112-
end
113-
11491
function extract_jacobian!(::Type{T}, result::AbstractArray, ydual::AbstractArray, n) where {T}
11592
out_reshaped = reshape(result, length(ydual), n)
11693
ydual_reshaped = vec(ydual)
@@ -180,27 +157,6 @@ function vector_mode_jacobian!(result, f!::F, y, x, cfg::JacobianConfig{T}) wher
180157
return result
181158
end
182159

183-
@inline function vector_mode_jacobian(f, x::StaticArray)
184-
T = typeof(Tag(f, eltype(x)))
185-
return extract_jacobian(T, static_dual_eval(T, f, x), x)
186-
end
187-
188-
@inline function vector_mode_jacobian!(result, f, x::StaticArray)
189-
T = typeof(Tag(f, eltype(x)))
190-
ydual = static_dual_eval(T, f, x)
191-
result = extract_jacobian!(T, result, ydual, length(x))
192-
result = extract_value!(T, result, ydual)
193-
return result
194-
end
195-
196-
@inline function vector_mode_jacobian!(result::ImmutableDiffResult, f, x::StaticArray)
197-
T = typeof(Tag(f, eltype(x)))
198-
ydual = static_dual_eval(T, f, x)
199-
result = DiffResults.jacobian!(result, extract_jacobian(T, ydual, x))
200-
result = DiffResults.value!(d -> value(T,d), result, ydual)
201-
return result
202-
end
203-
204160
const JACOBIAN_ERROR = DimensionMismatch("jacobian(f, x) expects that f(x) is an array. Perhaps you meant gradient(f, x)?")
205161

206162
# chunk mode #

0 commit comments

Comments
 (0)