Fix eigvals gradient for symmetric matrices (FluxML/Zygote#1369)

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRules"
 uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2"
-version = "1.73.0"
+version = "1.73.1"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/rulesets/LinearAlgebra/factorization.jl

@@ -335,8 +335,26 @@ function rrule(::typeof(eigen), A::StridedMatrix{T}; kwargs...) where {T<:Union{
             hermA = Hermitian(A)
             ∂V = ΔV isa AbstractZero ? ΔV : copyto!(similar(ΔV), ΔV)
             ∂hermA = eigen_rev!(hermA, λ, V, Δλ, ∂V)
-            ∂Atriu = _symherm_back(typeof(hermA), ∂hermA, Symbol(hermA.uplo))
-            ∂A = ∂Atriu isa AbstractTriangular ? triu!(∂Atriu.data) : ∂Atriu
+            if T <: Real && ΔV isa AbstractZero
+                # `A` is a plain matrix whose entries are all independent free
+                # variables; it merely happens to be symmetric, so the primal
+                # dispatched to the symmetric algorithm. `eigen_rev!` already returns
+                # a (Symmetric-wrapped) cotangent with the off-diagonal eigenvalue
+                # sensitivity split evenly across both triangles. Materialise that
+                # full matrix rather than projecting onto the stored triangle: the
+                # projection zeroes out the other triangle, disagreeing with
+                # ForwardDiff/finite differences and producing a gradient that is
+                # discontinuous as `A` crosses exact symmetry (see #1369).
+                #
+                # This is limited to the eigenvalues of a real matrix: the
+                # eigenvector phase convention differs between the symmetric and
+                # general algorithms, and the complex-Hermitian case is hard to pin
+                # down against finite differences, so those paths are left as-is.
+                ∂A = ∂hermA isa AbstractZero ? ∂hermA : Matrix(∂hermA)
+            else
+                ∂Atriu = _symherm_back(typeof(hermA), ∂hermA, Symbol(hermA.uplo))
+                ∂A = ∂Atriu isa AbstractTriangular ? triu!(∂Atriu.data) : ∂Atriu
+            end
         elseif ΔV isa AbstractZero
             ∂K = Diagonal(Δλ)
             ∂A = V' \ ∂K * V'

test/rulesets/LinearAlgebra/factorization.jl

@@ -349,8 +349,13 @@ end
                         ∂F_stable = (; [s => copy(getproperty(ΔF, s)) for s in nzprops]...)
                         :vectors in nzprops && rmul!(∂F_stable.vectors, C)
 
+                        # For the eigenvalues of a real matrix, the cotangent is the
+                        # general-matrix gradient (matching the unwrapped `eigen`); for
+                        # everything else it follows the symmetric-manifold convention
+                        # (FD through `Matrix(Hermitian(x))`). See #1369.
+                        wrap = (T <: Real && nzprops == [:values]) ? identity : x -> Matrix(Hermitian(x))
                         f_stable = function(x)
-                            F_ = _eigen_stable(Matrix(Hermitian(x)))
+                            F_ = _eigen_stable(wrap(x))
                             return (; (s => getproperty(F_, s) for s in nzprops)...)
                         end
 
@@ -414,7 +419,18 @@ end
                     ∂self, ∂A = @maybe_inferred back(Δλ)
                     @test ∂self === NoTangent()
                     @test ∂A isa typeof(A)
-                    @test ∂A ≈ j′vp(_fdm, A -> eigvals(Matrix(Hermitian(A))), Δλ, A)[1]
+                    if T <: Real
+                        # `A` is a plain matrix that merely happens to be symmetric,
+                        # so the cotangent must be the general-matrix gradient (split
+                        # over both triangles), matching the unwrapped `eigvals`, not
+                        # the symmetric-manifold gradient projected to one triangle
+                        # (see #1369). FiniteDifferences cannot diff the unwrapped
+                        # complex `eigvals` (eigenvalues leave the reals), so the
+                        # complex case keeps the Hermitian-wrapped reference.
+                        @test ∂A ≈ j′vp(_fdm, eigvals, Δλ, A)[1]
+                    else
+                        @test ∂A ≈ j′vp(_fdm, A -> eigvals(Matrix(Hermitian(A))), Δλ, A)[1]
+                    end
                     @test @maybe_inferred(back(ZeroTangent())) == (NoTangent(), ZeroTangent())
                 end
             end