Use Functor for ProjectTo. Update test

Author: Alex Robson

src/chainrules.jl

@@ -3,30 +3,25 @@
 function ChainRulesCore.rrule(::typeof(*), A::Real, B::WoodburyPDMat)
     project_A = ProjectTo(A)
     project_B = ProjectTo(B)
-    function times_pullback(Ȳ::AbstractMatrix)
-        Ā = @thunk(project_A(dot(Ȳ, B)'))
-        B̄ = @thunk(project_B(A' * Ȳ))
-        return (NoTangent(), Ā, B̄)
-    end
-
-    function times_pullback(Ȳ::Tangent{<:WoodburyPDMat})
-        Ā = dot(Ȳ.A * Ȳ.D * Ȳ.A' + Ȳ.S, B)
-        B̄ = Ȳ.A * (A' * Ȳ.D) * Ȳ.A' + A' * Ȳ.S
-        return (
-           NoTangent(),
-           @thunk(project_A(Ā')),
-           @thunk(project_B(B̄)),
-        )
-    end
+    times_pullback(ȳ) = _times_pullback(ȳ, A, B, (;A=project_A, B=project_B))
     return A * B, times_pullback
 end
 
-function ChainRulesCore.ProjectTo(W::WoodburyPDMat)
-    function dW(W̄)
-        Ā(W̄) = ProjectTo(W.A)(collect((W.D * W.A' * W̄' + W.D * W.A' * W̄)'))
-        D̄(W̄) = ProjectTo(W.D)(W.A' * (W̄) * W.A)
-        S̄(W̄) = ProjectTo(W.S)(W̄)
-        return Tangent{typeof(W)}(; A = Ā(W̄), D = D̄(W̄), S = S̄(W̄))
-    end
-    return dW
+function _times_pullback(Ȳ::AbstractMatrix, A, B, proj)
+    Ā = proj.A(dot(Ȳ, B)')
+    B̄ = proj.B(A' * Ȳ)
+    return (NoTangent(), Ā, B̄)
+end
+_times_pullback(ȳ::AbstractThunk, A, B, proj) = _times_pullback(unthunk(ȳ), A, B, proj)
+
+function ChainRulesCore.ProjectTo(W::T) where {T<:WoodburyPDMat}
+    fields = (A = W.A, D = W.D, S = W.S)
+    ProjectTo{T}(; fields...)
+end
+
+function (W::ProjectTo{T})(W̄) where {T<:WoodburyPDMat}
+    Ā(W̄) = ProjectTo(W.A)((W̄ + W̄') * (W.A * W.D))
+    D̄(W̄) = ProjectTo(W.D)(W.A' * (W̄) * W.A)
+    S̄(W̄) = ProjectTo(W.S)(W̄)
+    return Tangent{T}(; A = Ā(W̄), D = D̄(W̄), S = S̄(W̄))
 end

test/chainrules.jl

@@ -71,6 +71,27 @@ LinearAlgebra.dot(A, B::WoodburyLike) = dot(A, WoodburyPDMat(B))
             # Cannot get this to work. Here the T will be 
             # T = rand_tangent(primal::WoodburyPDMat) which breaks. 
             # test_rrule(*, 5.0, W)
+
+            #####################################################################################################################
+
+            primal = R * W
+
+            # Generate the Tangent as ChainRulesTestUtils would do
+            ∂primal = rand_tangent(Random.GLOBAL_RNG, collect(primal))
+            T = ProjectTo(primal)(∂primal)
+      
+            f_jvp = j′vp(ChainRulesTestUtils._fdm, x -> (*(x...)), T, (R, W))[1]
+
+            # Expected
+            R̄ = ProjectTo(R)(dot(∂primal, W'))
+            W̄ = ProjectTo(W)(conj(R) * ∂primal)
+
+            @test res[1] == primal
+            @test R̄ ≈ f_jvp[1]
+            @test W̄.A ≈ f_jvp[2].A
+            @test W̄.D ≈ f_jvp[2].D
+            @test W̄.S ≈ f_jvp[2].S
+
         end
     end
 end