Diag tests working for Mooncake

Author: kshyatt

ext/MatrixAlgebraKitMooncakeExt/MatrixAlgebraKitMooncakeExt.jl

@@ -40,18 +40,24 @@ for (f!, f, pb, adj) in (
             arg2c = copy(arg2)
             $f!(A, args, Mooncake.primal(alg_dalg))
             function $adj(::NoRData)
-                copy!(A, Ac)
+                # DON'T copy Ac to A if A === one 
+                # of the output args -- this can
+                # mess up the pullback because
+                # generally the args are used there
                 if !(A === arg1 || A === arg2)
+                    copy!(A, Ac)
                     $pb(dA, A, (arg1, arg2), (darg1, darg2))
                 else
                     ΔA = zero(A)
                     $pb(ΔA, A, (arg1, arg2), (darg1, darg2))
                     dA .= ΔA
                 end
                 if A === arg1
+                    copy!(A, Ac)
                     zero!(darg2)
                     copy!(arg2, arg2c)
                 elseif A === arg2
+                    copy!(A, Ac)
                     zero!(darg1)
                     copy!(arg1, arg1c)
                 else
@@ -60,7 +66,7 @@ for (f!, f, pb, adj) in (
                     copy!(arg2, arg2c)
                     copy!(arg1, arg1c)
                 end
-                return NoRData(), NoRData(), NoRData(), NoRData()
+                return ntuple(Returns(NoRData()), 4)
             end
             return args_dargs, $adj
         end

test/mooncake/eig.jl

@@ -16,6 +16,6 @@ for T in (BLASFloats..., GenericFloats...)
     if !is_buildkite
         TestSuite.test_mooncake_eig(T, (m, m); atol = m * m * TestSuite.precision(T), rtol = m * m * TestSuite.precision(T))
         AT = Diagonal{T, Vector{T}}
-        TestSuite.test_mooncake_eig(AT, m; atol = m * m * TestSuite.precision(T), rtol = m * m * TestSuite.precision(T))
+        TestSuite.test_mooncake_eig(AT, (m, m); atol = m * m * TestSuite.precision(T), rtol = m * m * TestSuite.precision(T))
     end
 end

test/mooncake/lq.jl

@@ -15,9 +15,9 @@ for T in (BLASFloats..., GenericFloats...), n in (17, m, 23)
     TestSuite.seed_rng!(1234)
     if !is_buildkite
         TestSuite.test_mooncake_lq(T, (m, n); atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))
-        #=if m == n
+        if m == n
             AT = Diagonal{T, Vector{T}}
-            TestSuite.test_mooncake_lq(AT, m; atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))
-        end=# # broken with singular exception
+            TestSuite.test_mooncake_lq(AT, (m, m); atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))
+        end
     end
 end

test/mooncake/orthnull.jl

@@ -15,9 +15,9 @@ for T in (BLASFloats..., GenericFloats...), n in (17, m, 23)
     TestSuite.seed_rng!(1234)
     if !is_buildkite
         TestSuite.test_mooncake_orthnull(T, (m, n); atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))
-        #=if m == n
+        if m == n
             AT = Diagonal{T, Vector{T}}
             TestSuite.test_mooncake_orthnull(AT, m; atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))
-        end=#
+        end
     end
 end

test/mooncake/qr.jl

@@ -15,9 +15,9 @@ for T in (BLASFloats..., GenericFloats...), n in (17, m, 23)
     TestSuite.seed_rng!(1234)
     if !is_buildkite
         TestSuite.test_mooncake_qr(T, (m, n); atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))
-        #=if m == n
+        if m == n
             AT = Diagonal{T, Vector{T}}
-            TestSuite.test_mooncake_qr(AT, m; atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))
-        end=#
+            TestSuite.test_mooncake_qr(AT, (m, m); atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))
+        end
     end
 end

test/testsuite/TestSuite.jl

@@ -96,6 +96,12 @@ function instantiate_rank_deficient_matrix(T, sz; trunc = truncrank(div(min(sz..
     return mul!(A, V, C)
 end
 
+function instantiate_rank_deficient_matrix(::Type{T}, sz; trunc = truncrank(div(min(sz...), 2))) where {T <: Diagonal}
+    A = instantiate_matrix(eltype(T), sz)
+    V, C = left_orth!(A; trunc)
+    return Diagonal(diag(mul!(A, V, C)))
+end
+
 include("ad_utils.jl")
 
 include("projections.jl")