Tidy up tests. White space clear

Author: Alex Robson

Project.toml

@@ -1,12 +1,10 @@
 name = "PDMatsExtras"
 uuid = "2c7acb1b-7338-470f-b38f-951d2bcb9193"
 authors = ["Invenia Technical Computing"]
-version = "2.5.2"
+version = "2.5.1"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
-FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
@@ -18,10 +16,12 @@ ChainRulesCore = "1"
 Distributions = "0.23, 0.24"
 FiniteDifferences = "0.11, 0.12"
 PDMats = "0.9, 0.10, 0.11"
-Zygote = "0.4, 0.5, 0.6"
+Zygote = "0.6"
 julia = "1"
 
 [extras]
+ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
+ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -30,4 +30,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Distributions", "FiniteDifferences", "Random", "SuiteSparse", "Test", "Zygote"]
+test = ["ChainRulesTestUtils", "ChainRules", "Distributions", "FiniteDifferences", "Random", "SuiteSparse", "Test", "Zygote"]

src/chainrules.jl

@@ -5,7 +5,6 @@ function ChainRulesCore.rrule(
 ) where {T, TA, TD, TS}
     project_A = ProjectTo(A)
     project_B = ProjectTo(B)
-
     function times_pullback(ȳ::AbstractMatrix)
         Ȳ = unthunk(ȳ)
         Ā = dot(Ȳ, B)
@@ -38,9 +37,4 @@ function ChainRulesCore.ProjectTo(W::WoodburyPDMat)
         return Tangent{typeof(W)}(; A = Ā(W̄), D = D̄(W̄), S = S̄(W̄))
     end
     return dW
-end
-
-
-
-
-
+end

test/chainrules.jl

@@ -50,22 +50,27 @@ LinearAlgebra.dot(A, B::WoodburyLike) = dot(A, WoodburyPDMat(B))
             ###
 
             primal = R * W
-            
             # Matrix Tangent
             T = rand_tangent(Matrix(primal))
+            res = ChainRulesCore.rrule(*, R, W)
             f_jvp = j′vp(ChainRulesTestUtils._fdm, x -> Matrix(*(x...)), T, (R, W))[1]
-            R̄ = dot(T, W')
-            W̄ = conj(R) * T
 
-            R̄ ≈ f_jvp[1]
-            (W.D * W.A' * W̄' + W.D * W.A' * W̄) ≈ f_jvp[2].A' # A transpose.
-            Diagonal(W.A' * (W̄) * W.A) ≈ f_jvp[2].D # D
-            Diagonal(W̄) ≈ f_jvp[2].S # S
+            # Expected
+            R̄ = ProjectTo(R)(dot(T, W'))
+            W̄ = ProjectTo(W)(conj(R) * T)
+
+            R̄_rrule = unthunk(res[2](T)[2])
+            W̄_rrule = unthunk(res[2](T)[3])
+
+            @test res[1] == primal
+            @test R̄_rrule ≈ f_jvp[1]
+            @test W̄_rrule.A ≈ f_jvp[2].A
+            @test W̄_rrule.D ≈ f_jvp[2].D
+            @test W̄_rrule.S ≈ f_jvp[2].S
 
             # Cannot get this to work. Here the T will be 
             # T = rand_tangent(primal::WoodburyPDMat) which breaks. 
             # test_rrule(*, 5.0, W)
-
         end
     end
 end