Use natural tangents in Matrix(::HermOrSym) rules

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRules"
 uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2"
-version = "1.74.6"
+version = "1.74.7"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -18,7 +18,7 @@ StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
 
 [compat]
 Adapt = "3.4.0, 4"
-ChainRulesCore = "1.25"
+ChainRulesCore = "1.26.6"
 Distributed = "1"
 GPUArraysCore = "0.1.0, 0.2"
 IrrationalConstants = "0.1.1, 0.2"

src/rulesets/LinearAlgebra/symmetric.jl

@@ -16,36 +16,18 @@ function rrule(T::Type{<:LinearAlgebra.HermOrSym}, A::AbstractMatrix, uplo)
 end
 
 function frule((_, ΔA), TM::Type{<:Matrix}, A::LinearAlgebra.HermOrSym)
-    return TM(A), TM(_symherm_forward(A, ΔA))
-end
-function frule((_, ΔA), ::Type{Array}, A::LinearAlgebra.HermOrSym)
-    return Array(A), Array(_symherm_forward(A, ΔA))
+    ∂A = ProjectTo(A)(unthunk(ΔA))
+    return TM(A), ∂A isa AbstractZero ? ∂A : TM(∂A)
 end
+frule(Δargs, ::Type{Array}, A::LinearAlgebra.HermOrSym) = frule(Δargs, Matrix, A)
 
 function rrule(TM::Type{<:Matrix}, A::LinearAlgebra.HermOrSym)
-    function Matrix_pullback(Ω̄)
-        ΔΩ = unthunk(Ω̄)
-        TA = _symhermtype(A)
-        T∂A = TA{eltype(ΔΩ),typeof(ΔΩ)}
-        uplo = A.uplo
-        ∂A = T∂A(_symherm_back(typeof(A), ΔΩ, Symbol(uplo)), uplo)
-        return NoTangent(), ∂A
-    end
+    project_A = ProjectTo(A)
+    Matrix_pullback(ΔΩ) = (NoTangent(), project_A(ΔΩ))
     return TM(A), Matrix_pullback
 end
 rrule(::Type{Array}, A::LinearAlgebra.HermOrSym) = rrule(Matrix, A)
 
-# for Ω = Matrix(A::HermOrSym), push forward ΔA to get ∂Ω
-function _symherm_forward(A, Ȧ)
-    ΔA = unthunk(Ȧ)
-    TA = _symhermtype(A)
-    return if typeof(ΔA) <: TA
-        ΔA
-    else
-        TA{eltype(ΔA),typeof(ΔA)}(ΔA, A.uplo)
-    end
-end
-
 # for Ω = HermOrSym(A, uplo), pull back ΔΩ to get ∂A
 @inline function _symherm_back(::Type{T}, Ω̄, uplo::Symbol) where {T}
     ΔΩ = unthunk(Ω̄)

test/Project.toml

@@ -16,9 +16,9 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]
 Adapt = "3.4.0, 4"
-ChainRulesCore = "1.25"
+ChainRulesCore = "1.26.6"
 ChainRulesTestUtils = "1.5"
-FiniteDifferences = "0.12.20"
+FiniteDifferences = "0.12.35"
 GPUArraysCore = "0.1.0, 0.2"
 JLArrays = "0.1, 0.4"
 LinearAlgebra = "1"

test/rulesets/Base/indexing.jl

@@ -102,6 +102,7 @@ end
         @test unthunk(dgrad) ≈ Diagonal([0, 2, 0])
 
         test_rrule(getindex, Symmetric(rand(3, 3)), 2, 2; check_inferred=false)  # Infers to Any
+        test_rrule(getindex, Symmetric(rand(3, 3)), 2, 3; check_inferred=false)
         sgrad = rrule(getindex, Symmetric(rand(3, 3)), 2, 3)[2](1.0)[2]
         @test unthunk(sgrad) ≈ [0 0 0; 0 0 1/2; 0 1/2 0]
     end

test/rulesets/Base/mapreduce.jl

@@ -208,9 +208,8 @@ struct SumRuleConfig <: RuleConfig{Union{HasReverseMode}} end
                 # Symmetric -- at least this doesn't have zeros, still an unlikely combination
                 xs = Symmetric(rand(T,4,4))
                 @test unthunk(rrule(prod, Symmetric(T[1 2; -333 4]))[2](1.0)[2]) == [16 8; 8 4]
-                # TODO debug why these fail  https://github.com/JuliaDiff/ChainRules.jl/issues/475
-                @test_skip test_rrule(prod, xs)
-                @test_skip test_rrule(prod, xs, fkwargs=(dims=2,))
+                test_rrule(prod, xs)
+                test_rrule(prod, xs, fkwargs=(dims=2,))
             end
         end
         @testset "Array{Float32}, no zero entries" begin

test/rulesets/LinearAlgebra/symmetric.jl

@@ -56,6 +56,13 @@ using ChainRulesTestUtils: _fdm
         # intentionally specifying tangents here to test both SymHerm (default) and Matrix
         test_frule(f, x)
         test_frule(f, x ⊢ randn(T, 3, 3))
+        @test frule((NoTangent(), ZeroTangent()), f, x) == (f(x), ZeroTangent())
+
+        Δx = SymHerm(randn(T, 3, 3), uplo === :U ? :L : :U)
+        @test frule((NoTangent(), Δx), f, x)[2] == Δx
+        ∂x = rrule(f, x)[2](Δx)[2]
+        @test ∂x == Δx
+        @test ∂x.uplo == x.uplo
     end
 
     # symmetric/hermitian eigendecomposition follows the sign convention