Turn StaticArrays into an Extension Package

Project.toml

@@ -28,12 +28,19 @@ SpecialFunctions = "0.8, 0.9, 0.10, 1.0, 2"
 StaticArrays = "0.8.3, 0.9, 0.10, 0.11, 0.12, 1.0"
 julia = "1"
 
+[extensions]
+ForwardDiffStaticArraysExt = "StaticArrays"
+
 [extras]
 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", "Test", "InteractiveUtils", "StaticArrays"]
+
+[weakdeps]
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"

ext/ForwardDiffStaticArraysExt.jl

@@ -0,0 +1,136 @@
+module ForwardDiffStaticArraysExt
+
+using ForwardDiff, StaticArrays, LinearAlgebra, 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*ForwardDiff._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
+
+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(f, x::StaticArray)
+    T = typeof(Tag(f, eltype(x)))
+    return extract_jacobian(T, static_dual_eval(T, f, x), 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

src/ForwardDiff.jl

@@ -2,7 +2,6 @@ module ForwardDiff
 
 using DiffRules, DiffResults
 using DiffResults: DiffResult, MutableDiffResult, ImmutableDiffResult
-using StaticArrays
 if VERSION >= v"1.6"
     using Preferences
 end
@@ -25,6 +24,10 @@ include("gradient.jl")
 include("jacobian.jl")
 include("hessian.jl")
 
+if !isdefined(Base, :get_extension)
+    include("../ext/ForwardDiffStaticArraysExt.jl")
+end
+
 export DiffResults
 
 end # module

src/apiutils.jl

@@ -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)

src/dual.jl

@@ -708,12 +708,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)
@@ -743,13 +737,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))))

src/gradient.jl

@@ -41,29 +41,12 @@ 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!(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)
-
 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)
@@ -115,16 +98,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 #
 ##############

src/hessian.jl

@@ -67,28 +67,3 @@ function hessian!(result::DiffResult, f, x::AbstractArray, cfg::HessianConfig{T}
     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!(result::AbstractArray, f, x::StaticArray) = jacobian!(result, y -> gradient(f, y), x)
-
-hessian!(result::MutableDiffResult, f, x::StaticArray) = 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)
-
-function 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

src/jacobian.jl

@@ -82,35 +82,12 @@ 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!(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)
-
 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)
@@ -180,27 +157,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 #