Add support for StaticArrays >= 1.7 (#703)

Author: devmotion

Project.toml

@@ -32,7 +32,7 @@ LogExpFunctions = "0.3"
 NaNMath = "1"
 Preferences = "1"
 SpecialFunctions = "1, 2"
-StaticArrays = "1.5 - 1.6"
+StaticArrays = "1.5"
 julia = "1.6"
 
 [extras]

ext/ForwardDiffStaticArraysExt.jl

@@ -7,7 +7,7 @@ using ForwardDiff: Dual, partials, GradientConfig, JacobianConfig, HessianConfig
                    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!
+                   vector_mode_jacobian, vector_mode_jacobian!, valtype, value
 using DiffResults: DiffResult, ImmutableDiffResult, MutableDiffResult
 
 @generated function dualize(::Type{T}, x::StaticArray) where T
@@ -23,19 +23,17 @@ end
 
 @inline static_dual_eval(::Type{T}, f, x::StaticArray) where T = f(dualize(T, x))
 
+# To fix method ambiguity issues:
 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...))
+    return ForwardDiff._eigvals(A)
 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...)))
+    return ForwardDiff._eigen(A)
 end
 
+# For `MMatrix` we can use the in-place method
+ForwardDiff._lyap_div!!(A::StaticArrays.MMatrix, λ::AbstractVector) = ForwardDiff._lyap_div!(A, λ)
+
 # Gradient
 @inline ForwardDiff.gradient(f, x::StaticArray)                      = vector_mode_gradient(f, x)
 @inline ForwardDiff.gradient(f, x::StaticArray, cfg::GradientConfig) = gradient(f, x)

src/dual.jl

@@ -719,7 +719,11 @@ end
 # Symmetric eigvals #
 #-------------------#
 
-function LinearAlgebra.eigvals(A::Symmetric{<:Dual{Tg,T,N}}) where {Tg,T<:Real,N}
+# To be able to reuse this default definition in the StaticArrays extension
+# (has to be re-defined to avoid method ambiguity issues)
+# we forward the call to an internal method that can be shared and reused
+LinearAlgebra.eigvals(A::Symmetric{<:Dual{Tg,T,N}}) where {Tg,T<:Real,N} = _eigvals(A)
+function _eigvals(A::Symmetric{<:Dual{Tg,T,N}}) 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...))
@@ -737,8 +741,19 @@ function LinearAlgebra.eigvals(A::SymTridiagonal{<:Dual{Tg,T,N}}) where {Tg,T<:R
     Dual{Tg}.(λ, tuple.(parts...))
 end
 
-# A ./ (λ - λ') but with diag special cased
-function _lyap_div!(A, λ)
+# A ./ (λ' .- λ) but with diag special cased
+# Default out-of-place method
+function _lyap_div!!(A::AbstractMatrix, λ::AbstractVector)
+    return map(
+        (a, b, idx) -> a / (idx[1] == idx[2] ? oneunit(b) : b),
+        A,
+        λ' .- λ,
+        CartesianIndices(A),
+    )
+end
+# For `Matrix` (and e.g. `StaticArrays.MMatrix`) we can use an in-place method
+_lyap_div!!(A::Matrix, λ::AbstractVector) = _lyap_div!(A, λ)
+function _lyap_div!(A::AbstractMatrix, λ::AbstractVector)
     for (j,μ) in enumerate(λ), (k,λ) in enumerate(λ)
         if k ≠ j
             A[k,j] /= μ - λ
@@ -747,17 +762,21 @@ function _lyap_div!(A, λ)
     A
 end
 
-function LinearAlgebra.eigen(A::Symmetric{<:Dual{Tg,T,N}}) where {Tg,T<:Real,N}
+# To be able to reuse this default definition in the StaticArrays extension
+# (has to be re-defined to avoid method ambiguity issues)
+# we forward the call to an internal method that can be shared and reused
+LinearAlgebra.eigen(A::Symmetric{<:Dual{Tg,T,N}}) where {Tg,T<:Real,N} = _eigen(A)
+function _eigen(A::Symmetric{<:Dual{Tg,T,N}}) 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)
+    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))))
-    parts = ntuple(j -> Q*_lyap_div!(Q' * getindex.(partials.(A), j) * Q - Diagonal(getindex.(partials.(λ), j)), value.(λ)), N)
+    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
 

test/JacobianTest.jl

@@ -243,8 +243,14 @@ end
     @test ForwardDiff.jacobian(ev1, x0) ≈ Calculus.finite_difference_jacobian(ev1, x0)
     ev2(x) = eigen(SymTridiagonal(x, x[1:end-1])).vectors[:,1]
     @test ForwardDiff.jacobian(ev2, x0) ≈ Calculus.finite_difference_jacobian(ev2, x0)
-    x0_static = SVector{2}(x0)
-    @test ForwardDiff.jacobian(ev1, x0_static) ≈ Calculus.finite_difference_jacobian(ev1, x0)
+
+    x0_svector = SVector{2}(x0)
+    @test ForwardDiff.jacobian(ev1, x0_svector) isa SMatrix{2, 2}
+    @test ForwardDiff.jacobian(ev1, x0_svector) ≈ Calculus.finite_difference_jacobian(ev1, x0)
+
+    x0_mvector = MVector{2}(x0)
+    @test ForwardDiff.jacobian(ev1, x0_mvector) isa MMatrix{2, 2}
+    @test ForwardDiff.jacobian(ev1, x0_mvector) ≈ Calculus.finite_difference_jacobian(ev1, x0)
 end
 
 @testset "type stability" begin

test/MiscTest.jl

@@ -6,7 +6,9 @@ using Test
 using ForwardDiff
 using DiffTests
 using SparseArrays: sparse
+using StaticArrays
 using IrrationalConstants
+using LinearAlgebra
 
 include(joinpath(dirname(@__FILE__), "utils.jl"))
 
@@ -180,4 +182,41 @@ end
 # example from https://github.com/JuliaDiff/DiffRules.jl/pull/98#issuecomment-1574420052
 @test only(ForwardDiff.hessian(t -> abs(t[1])^2, [0.0])) == 2
 
+@testset "_lyap_div!!" begin
+    # In-place version for `Matrix`
+    A = rand(3, 3)
+    Acopy = copy(A)
+    λ = rand(3)
+    B = @inferred(ForwardDiff._lyap_div!!(A, λ))
+    @test B === A
+    @test B[diagind(B)] == Acopy[diagind(Acopy)]
+    no_diag(X) = [X[i] for i in eachindex(X) if !(i in diagind(X))]
+    @test no_diag(B) == no_diag(Acopy ./ (λ' .- λ))
+
+    # Immutable static arrays
+    A_smatrix = SMatrix{3,3}(Acopy)
+    λ_svector = SVector{3}(λ)
+    B_smatrix = @inferred(ForwardDiff._lyap_div!!(A_smatrix, λ_svector))
+    @test B_smatrix !== A_smatrix
+    @test B_smatrix isa SMatrix{3,3}
+    @test B_smatrix == B
+    λ_mvector = MVector{3}(λ)
+    B_smatrix = @inferred(ForwardDiff._lyap_div!!(A_smatrix, λ_mvector))
+    @test B_smatrix !== A_smatrix
+    @test B_smatrix isa SMatrix{3,3}
+    @test B_smatrix == B
+
+    # Mutable static arrays
+    A_mmatrix = MMatrix{3,3}(Acopy)
+    λ_svector = SVector{3}(λ)
+    B_mmatrix = @inferred(ForwardDiff._lyap_div!!(A_mmatrix, λ_svector))
+    @test B_mmatrix === A_mmatrix
+    @test B_mmatrix == B
+    A_mmatrix = MMatrix{3,3}(Acopy)
+    λ_mvector = MVector{3}(λ)
+    B_mmatrix = @inferred(ForwardDiff._lyap_div!!(A_mmatrix, λ_mvector))
+    @test B_mmatrix === A_mmatrix
+    @test B_mmatrix == B
+end
+
 end # module