Mooncake testsuite refactor (#175)

* small refactor QR pullback

add QR gauge projection

* testsuite reorganisation

* add QR mooncake tests

* Genius suggestion by @Jutho fixes everything

* Refactor Mooncake LQ tests

* Refactor Mooncake Eig tests

* fix pullback implementations!

* Refactor Mooncake SVD tests

* Refactor Mooncake Polar tests

* make testsets verbose

* Refactor Mooncake OrthNull tests

* clean up

* rename `call_and_zero!`

* move gauge dependence removal to ad_utils again

* separate out mooncake tests

* Update src/pullbacks/qr.jl

Co-authored-by: Jutho <Jutho@users.noreply.github.com>

* move functions to ad_utils

* more ad_setup usage

* add back rank_deficient tests

* clean up QR/LQ pullbacks

* add back `trunctol` tests

* more more ad_setup usage + chainrules simplification

* make qr/lq gauge dependence also fix R/L

---------

Co-authored-by: Jutho <Jutho@users.noreply.github.com>

Author: lkdvos

ext/MatrixAlgebraKitMooncakeExt/MatrixAlgebraKitMooncakeExt.jl

@@ -239,7 +239,7 @@ for f in (:eig, :eigh)
                 _warn_pullback_truncerror(dϵ)
 
                 # compute pullbacks
-                $f_pullback!(dA, Ac, DVc, dDVtrunc, ind)
+                $f_pullback!(dA, Ac, DV, dDVtrunc, ind)
                 zero!.(dDVtrunc) # since this is allocated in this function this is probably not required
 
                 # restore state
@@ -351,8 +351,8 @@ for f in (:eig, :eigh)
             dDVtrunc = last.(arrayify.(DVtrunc, Mooncake.tangent(DVtrunc_dDVtrunc)))
             function $f_adjoint!(::NoRData)
                 # compute pullbacks
-                $f_pullback!(dA, Ac, DVc, dDVtrunc, ind)
-                zero!.(dDVtrunc) # since this is allocated in this function this is probably not required
+                $f_pullback!(dA, Ac, DV, dDVtrunc, ind)
+                zero!.(dDV)
 
                 # restore state
                 copy!(A, Ac)
@@ -425,7 +425,7 @@ for (f!, f) in (
             S, dS = arrayify(USVᴴ[2], dUSVᴴ[2])
             Vᴴ, dVᴴ = arrayify(USVᴴ[3], dUSVᴴ[3])
             USVᴴc = copy.(USVᴴ)
-            output = $f!(A, Mooncake.primal(alg_dalg))
+            output = $f!(A, USVᴴ, Mooncake.primal(alg_dalg))
             function svd_adjoint(::NoRData)
                 copy!(A, Ac)
                 if $(f! == svd_compact!)
@@ -590,7 +590,7 @@ function Mooncake.rrule!!(::CoDual{typeof(svd_trunc!)}, A_dA::CoDual, USVᴴ_dUS
         _warn_pullback_truncerror(dϵ)
 
         # compute pullbacks
-        svd_pullback!(dA, Ac, USVᴴc, dUSVᴴtrunc, ind)
+        svd_pullback!(dA, Ac, USVᴴ, dUSVᴴtrunc, ind)
         zero!.(dUSVᴴtrunc) # since this is allocated in this function this is probably not required
         zero!.(dUSVᴴ)
 
@@ -717,8 +717,7 @@ function Mooncake.rrule!!(::CoDual{typeof(svd_trunc_no_error!)}, A_dA::CoDual, U
     dUSVᴴtrunc = last.(arrayify.(USVᴴtrunc, Mooncake.tangent(USVᴴtrunc_dUSVᴴtrunc)))
     function svd_trunc_adjoint(::NoRData)
         # compute pullbacks
-        svd_pullback!(dA, Ac, USVᴴc, dUSVᴴtrunc, ind)
-        zero!.(dUSVᴴtrunc) # since this is allocated in this function this is probably not required
+        svd_pullback!(dA, Ac, USVᴴ, dUSVᴴtrunc, ind)
         zero!.(dUSVᴴ)
 
         # restore state

src/pullbacks/lq.jl

@@ -1,28 +1,33 @@
+lq_rank(L; kwargs...) = qr_rank(L; kwargs...)
+
 function check_lq_cotangents(
-        L, Q, ΔL, ΔQ, minmn::Int, p::Int;
+        L, Q, ΔL, ΔQ, p::Int;
         gauge_atol::Real = default_pullback_gauge_atol(ΔQ)
     )
+    minmn = min(size(L, 1), size(Q, 2))
     if minmn > p # case where A is rank-deficient
         Δgauge = abs(zero(eltype(Q)))
         if !iszerotangent(ΔQ)
             # in this case the number Householder reflections will
             # change upon small variations, and all of the remaining
-            # columns of ΔQ should be zero for a gauge-invariant
+            # rows of ΔQ should be zero for a gauge-invariant
             # cost function
             ΔQ2 = view(ΔQ, (p + 1):size(Q, 1), :)
-            Δgauge = max(Δgauge, norm(ΔQ2))
+            Δgauge_Q = norm(ΔQ2, Inf)
+            Δgauge = max(Δgauge, Δgauge_Q)
         end
         if !iszerotangent(ΔL)
             ΔL22 = view(ΔL, (p + 1):size(L, 1), (p + 1):minmn)
-            Δgauge = max(Δgauge, norm(ΔL22))
+            Δgauge_L = norm(ΔL22, Inf)
+            Δgauge = max(Δgauge, Δgauge_L)
         end
         Δgauge ≤ gauge_atol ||
             @warn "`lq` cotangents sensitive to gauge choice: (|Δgauge| = $Δgauge)"
     end
     return
 end
 
-function check_lq_full_cotangents(Q1, ΔQ2, ΔQ2Q1ᴴ; gauge_atol::Real = default_pullback_gauge_atol(Q1))
+function check_lq_full_cotangents(Q1, ΔQ2, ΔQ2Q1ᴴ; gauge_atol::Real = default_pullback_gauge_atol(ΔQ2))
     # in the case where A is full rank, but there are more columns in Q than in A
     # (the case of `lq_full`), there is gauge-invariant information in the
     # projection of ΔQ2 onto the column space of Q1, by virtue of Q being a unitary
@@ -32,7 +37,7 @@ function check_lq_full_cotangents(Q1, ΔQ2, ΔQ2Q1ᴴ; gauge_atol::Real = defaul
     # Q2' * ΔQ2 as a gauge dependent quantity.
     Δgauge = norm(mul!(copy(ΔQ2), ΔQ2Q1ᴴ, Q1, -1, 1), Inf)
     Δgauge ≤ gauge_atol ||
-        @warn "`lq` cotangents sensitive to gauge choice: (|Δgauge| = $Δgauge)"
+        @warn "`lq_full` cotangents sensitive to gauge choice: (|Δgauge| = $Δgauge)"
     return
 end
 
@@ -62,9 +67,7 @@ function lq_pullback!(
     L, Q = LQ
     m = size(L, 1)
     n = size(Q, 2)
-    minmn = min(m, n)
-    Ld = diagview(L)
-    p = @something findlast(>=(rank_atol) ∘ abs, Ld) 0
+    p = lq_rank(L; rank_atol)
 
     ΔL, ΔQ = ΔLQ
 
@@ -74,7 +77,7 @@ function lq_pullback!(
     ΔA1 = view(ΔA, 1:p, :)
     ΔA2 = view(ΔA, (p + 1):m, :)
 
-    check_lq_cotangents(L, Q, ΔL, ΔQ, minmn, p; gauge_atol)
+    check_lq_cotangents(L, Q, ΔL, ΔQ, p; gauge_atol)
 
     ΔQ̃ = zero!(similar(Q, (p, n)))
     if !iszerotangent(ΔQ)

src/pullbacks/qr.jl

@@ -1,4 +1,11 @@
-function check_qr_cotangents(Q, R, ΔQ, ΔR, minmn::Int, p::Int; gauge_atol::Real = default_pullback_gauge_atol(ΔQ))
+qr_rank(R; rank_atol = default_pullback_rank_atol(R)) =
+    @something findlast(>=(rank_atol) ∘ abs, diagview(R)) 0
+
+function check_qr_cotangents(
+        Q, R, ΔQ, ΔR, p::Int;
+        gauge_atol::Real = default_pullback_gauge_atol(ΔQ)
+    )
+    minmn = min(size(Q, 1), size(R, 2))
     if minmn > p # case where A is rank-deficient
         Δgauge = abs(zero(eltype(Q)))
         if !iszerotangent(ΔQ)
@@ -7,11 +14,13 @@ function check_qr_cotangents(Q, R, ΔQ, ΔR, minmn::Int, p::Int; gauge_atol::Rea
             # columns of ΔQ should be zero for a gauge-invariant
             # cost function
             ΔQ2 = view(ΔQ, :, (p + 1):size(Q, 2))
-            Δgauge = max(Δgauge, norm(ΔQ2, Inf))
+            Δgauge_Q = norm(ΔQ2, Inf)
+            Δgauge = max(Δgauge, Δgauge_Q)
         end
         if !iszerotangent(ΔR)
             ΔR22 = view(ΔR, (p + 1):minmn, (p + 1):size(R, 2))
-            Δgauge = max(Δgauge, norm(ΔR22, Inf))
+            Δgauge_R = norm(ΔR22, Inf)
+            Δgauge = max(Δgauge, Δgauge_R)
         end
         Δgauge ≤ gauge_atol ||
             @warn "`qr` cotangents sensitive to gauge choice: (|Δgauge| = $Δgauge)"
@@ -29,7 +38,7 @@ function check_qr_full_cotangents(Q1, ΔQ2, Q1dΔQ2; gauge_atol::Real = default_
     # Q2' * ΔQ2 as a gauge dependent quantity.
     Δgauge = norm(mul!(copy(ΔQ2), Q1, Q1dΔQ2, -1, 1), Inf)
     Δgauge ≤ gauge_atol ||
-        @warn "`qr` cotangents sensitive to gauge choice: (|Δgauge| = $Δgauge)"
+        @warn "`qr_full` cotangents sensitive to gauge choice: (|Δgauge| = $Δgauge)"
     return
 end
 
@@ -60,19 +69,17 @@ function qr_pullback!(
     Q, R = QR
     m = size(Q, 1)
     n = size(R, 2)
-    minmn = min(m, n)
     Rd = diagview(R)
-    p = @something findlast(>=(rank_atol) ∘ abs, Rd) 0
+    p = qr_rank(R; rank_atol)
 
     ΔQ, ΔR = ΔQR
 
     Q1 = view(Q, :, 1:p)
-    Q2 = view(Q, :, (p + 1):size(Q, 2))
     R11 = view(R, 1:p, 1:p)
     ΔA1 = view(ΔA, :, 1:p)
     ΔA2 = view(ΔA, :, (p + 1):n)
 
-    check_qr_cotangents(Q, R, ΔQ, ΔR, minmn, p; gauge_atol)
+    check_qr_cotangents(Q, R, ΔQ, ΔR, p; gauge_atol)
 
     ΔQ̃ = zero!(similar(Q, (m, p)))
     if !iszerotangent(ΔQ)

test/mooncake.jl

@@ -0,0 +1,19 @@
+using MatrixAlgebraKit
+using Test
+using LinearAlgebra: Diagonal
+using CUDA, AMDGPU
+
+BLASFloats = (Float32, ComplexF64) # full suite is too expensive on CI
+GenericFloats = ()
+@isdefined(TestSuite) || include("../testsuite/TestSuite.jl")
+using .TestSuite
+
+is_buildkite = get(ENV, "BUILDKITE", "false") == "true"
+
+m = 19
+for T in (BLASFloats..., GenericFloats...)
+    TestSuite.seed_rng!(123)
+    if !is_buildkite
+        TestSuite.test_mooncake_eig(T, (m, m); atol = m * m * TestSuite.precision(T), rtol = m * m * TestSuite.precision(T))
+    end
+end

test/mooncake/eig.jl

@@ -0,0 +1,19 @@
+using MatrixAlgebraKit
+using Test
+using LinearAlgebra: Diagonal
+using CUDA, AMDGPU
+
+BLASFloats = (Float32, ComplexF64) # full suite is too expensive on CI
+GenericFloats = ()
+@isdefined(TestSuite) || include("../testsuite/TestSuite.jl")
+using .TestSuite
+
+is_buildkite = get(ENV, "BUILDKITE", "false") == "true"
+
+m = 19
+for T in (BLASFloats..., GenericFloats...)
+    TestSuite.seed_rng!(123)
+    if !is_buildkite
+        TestSuite.test_mooncake_eigh(T, (m, m); atol = m * m * TestSuite.precision(T), rtol = m * m * TestSuite.precision(T))
+    end
+end

test/mooncake/eigh.jl

@@ -0,0 +1,19 @@
+using MatrixAlgebraKit
+using Test
+using LinearAlgebra: Diagonal
+using CUDA, AMDGPU
+
+BLASFloats = (Float32, ComplexF64) # full suite is too expensive on CI
+GenericFloats = ()
+@isdefined(TestSuite) || include("../testsuite/TestSuite.jl")
+using .TestSuite
+
+is_buildkite = get(ENV, "BUILDKITE", "false") == "true"
+
+m = 19
+for T in (BLASFloats..., GenericFloats...), n in (17, m, 23)
+    TestSuite.seed_rng!(123)
+    if !is_buildkite
+        TestSuite.test_mooncake_lq(T, (m, n); atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))
+    end
+end

test/mooncake/lq.jl

@@ -0,0 +1,19 @@
+using MatrixAlgebraKit
+using Test
+using LinearAlgebra: Diagonal
+using CUDA, AMDGPU
+
+BLASFloats = (Float32, ComplexF64) # full suite is too expensive on CI
+GenericFloats = ()
+@isdefined(TestSuite) || include("../testsuite/TestSuite.jl")
+using .TestSuite
+
+is_buildkite = get(ENV, "BUILDKITE", "false") == "true"
+
+m = 19
+for T in (BLASFloats..., GenericFloats...), n in (17, m, 23)
+    TestSuite.seed_rng!(123)
+    if !is_buildkite
+        TestSuite.test_mooncake_orthnull(T, (m, n); atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))
+    end
+end

test/mooncake/orthnull.jl

@@ -0,0 +1,21 @@
+using MatrixAlgebraKit
+using Test
+using LinearAlgebra: Diagonal
+using CUDA, AMDGPU
+
+BLASFloats = (Float32, ComplexF64) # full suite is too expensive on CI
+GenericFloats = ()
+@isdefined(TestSuite) || include("../testsuite/TestSuite.jl")
+using .TestSuite
+
+is_buildkite = get(ENV, "BUILDKITE", "false") == "true"
+
+m = 19
+for T in (BLASFloats..., GenericFloats...), n in (17, m, 23)
+    TestSuite.seed_rng!(123)
+    if !is_buildkite
+        atol = rtol = m * n * TestSuite.precision(T)
+        m >= n && TestSuite.test_mooncake_left_polar(T, (m, n); atol, rtol)
+        n >= m && TestSuite.test_mooncake_right_polar(T, (m, n); atol, rtol)
+    end
+end

test/mooncake/polar.jl

@@ -0,0 +1,19 @@
+using MatrixAlgebraKit
+using Test
+using LinearAlgebra: Diagonal
+using CUDA, AMDGPU
+
+BLASFloats = (Float32, ComplexF64) # full suite is too expensive on CI
+GenericFloats = ()
+@isdefined(TestSuite) || include("../testsuite/TestSuite.jl")
+using .TestSuite
+
+is_buildkite = get(ENV, "BUILDKITE", "false") == "true"
+
+m = 19
+for T in (BLASFloats..., GenericFloats...), n in (17, m, 23)
+    TestSuite.seed_rng!(123)
+    if !is_buildkite
+        TestSuite.test_mooncake_qr(T, (m, n); atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))
+    end
+end