Some changes to the ad test utils (#180)

* some changes to the ad test utils

* some more changes / lq tests are failing for unknown reasons

* one more attempt

* update enzyme

* unicode subscripts

* remove debug code

---------

Co-authored-by: Lukas Devos <ldevos98@gmail.com>

Author: Jutho

Project.toml

@@ -34,7 +34,7 @@ Enzyme = "0.13.118"
 EnzymeTestUtils = "0.2.5"
 GenericLinearAlgebra = "0.3.19"
 GenericSchur = "0.5.6"
-JET = "0.9, 0.10"
+JET = "0.9, 0.10, 0.11"
 LinearAlgebra = "1"
 Mooncake = "0.5"
 ParallelTestRunner = "2"

src/common/view.jl

@@ -24,7 +24,8 @@ diagonal(v::AbstractVector) = Diagonal(v)
 function lowertriangularind(A::AbstractMatrix)
     Base.require_one_based_indexing(A)
     m, n = size(A)
-    I = Vector{Int}(undef, div(m * (m - 1), 2) + m * (n - m))
+    minmn = min(m, n)
+    I = Vector{Int}(undef, div(minmn * (minmn - 1), 2) + minmn * (m - minmn))
     offset = 0
     for j in 1:n
         r = (j + 1):m
@@ -37,7 +38,8 @@ end
 function uppertriangularind(A::AbstractMatrix)
     Base.require_one_based_indexing(A)
     m, n = size(A)
-    I = Vector{Int}(undef, div(m * (m - 1), 2) + m * (n - m))
+    minmn = min(m, n)
+    I = Vector{Int}(undef, div(minmn * (minmn - 1), 2) + minmn * (n - minmn))
     offset = 0
     for i in 1:m
         r = (i + 1):n

src/pullbacks/lq.jl

@@ -5,40 +5,26 @@ function check_lq_cotangents(
         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
-            # rows of ΔQ should be zero for a gauge-invariant
-            # cost function
-            ΔQ2 = view(ΔQ, (p + 1):size(Q, 1), :)
-            Δ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_L = norm(ΔL22, Inf)
-            Δgauge = max(Δgauge, Δgauge_L)
-        end
-        Δgauge ≤ gauge_atol ||
-            @warn "`lq` cotangents sensitive to gauge choice: (|Δgauge| = $Δgauge)"
+    Δgauge = abs(zero(eltype(Q)))
+    if !iszerotangent(ΔQ)
+        ΔQ₂ = view(ΔQ, (p + 1):minmn, :)
+        ΔQ₃ = ΔQ[(minmn + 1):size(Q, 1), :]
+        Δgauge_Q = norm(ΔQ₂, Inf)
+        Q₁ = view(Q, 1:p, :)
+        ΔQ₃Q₁ᴴ = ΔQ₃ * Q₁'
+        mul!(ΔQ₃, ΔQ₃Q₁ᴴ, Q₁, -1, 1)
+        Δgauge_Q = max(Δgauge_Q, norm(ΔQ₃, Inf))
+        Δgauge = max(Δgauge, Δgauge_Q)
+    end
+    if !iszerotangent(ΔL)
+        ΔL22 = view(ΔL, (p + 1):size(ΔL, 1), (p + 1):minmn)
+        Δgauge_L = norm(view(ΔL22, lowertriangularind(ΔL22)), Inf)
+        Δgauge_L = max(Δgauge_L, norm(view(ΔL22, diagind(ΔL22)), Inf))
+        Δgauge = max(Δgauge, Δgauge_L)
     end
-    return
-end
-
-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
-    # matrix. As the number of Householder reflections is in fixed in the full rank
-    # case, Q is expected to rotate smoothly (we might even be able to predict) also
-    # how the full Q2 will change, but this we omit for now, and we consider
-    # Q2' * ΔQ2 as a gauge dependent quantity.
-    Δgauge = norm(mul!(copy(ΔQ2), ΔQ2Q1ᴴ, Q1, -1, 1), Inf)
     Δgauge ≤ gauge_atol ||
-        @warn "`lq_full` cotangents sensitive to gauge choice: (|Δgauge| = $Δgauge)"
-    return
+        @warn "`lq` cotangents sensitive to gauge choice: (|Δgauge| = $Δgauge)"
+    return nothing
 end
 
 """
@@ -67,54 +53,53 @@ function lq_pullback!(
     L, Q = LQ
     m = size(L, 1)
     n = size(Q, 2)
+    minmn = min(m, n)
     p = lq_rank(L; rank_atol)
 
     ΔL, ΔQ = ΔLQ
 
-    Q1 = view(Q, 1:p, :)
-    Q2 = view(Q, (p + 1):size(Q, 1), :)
-    L11 = view(L, 1:p, 1:p)
-    ΔA1 = view(ΔA, 1:p, :)
-    ΔA2 = view(ΔA, (p + 1):m, :)
+    Q₁ = view(Q, 1:p, :)
+    L₁₁ = LowerTriangular(view(L, 1:p, 1:p))
+    ΔA₁ = view(ΔA, 1:p, :)
+    ΔA₂ = view(ΔA, (p + 1):m, :)
 
     check_lq_cotangents(L, Q, ΔL, ΔQ, p; gauge_atol)
 
     ΔQ̃ = zero!(similar(Q, (p, n)))
     if !iszerotangent(ΔQ)
-        ΔQ1 = view(ΔQ, 1:p, :)
-        copy!(ΔQ̃, ΔQ1)
-        if p < size(Q, 1)
-            Q2 = view(Q, (p + 1):size(Q, 1), :)
-            ΔQ2 = view(ΔQ, (p + 1):size(Q, 1), :)
-            ΔQ2Q1ᴴ = ΔQ2 * Q1'
-            check_lq_full_cotangents(Q1, ΔQ2, ΔQ2Q1ᴴ; gauge_atol)
-            ΔQ̃ = mul!(ΔQ̃, ΔQ2Q1ᴴ', Q2, -1, 1)
+        ΔQ₁ = view(ΔQ, 1:p, :)
+        copy!(ΔQ̃, ΔQ₁)
+        if minmn < size(Q, 1)
+            ΔQ₃ = view(ΔQ, (minmn + 1):size(ΔQ, 1), :)
+            Q₃ = view(Q, (minmn + 1):size(Q, 1), :)
+            ΔQ₃Q₁ᴴ = ΔQ₃ * Q₁'
+            ΔQ̃ = mul!(ΔQ̃, ΔQ₃Q₁ᴴ', Q₃, -1, 1)
         end
     end
     if !iszerotangent(ΔL) && m > p
-        L21 = view(L, (p + 1):m, 1:p)
-        ΔL21 = view(ΔL, (p + 1):m, 1:p)
-        ΔQ̃ = mul!(ΔQ̃, L21' * ΔL21, Q1, -1, 1)
-        # Adding ΔA2 contribution
-        ΔA2 = mul!(ΔA2, ΔL21, Q1, 1, 1)
+        L₂₁ = view(L, (p + 1):m, 1:p)
+        ΔL₂₁ = view(ΔL, (p + 1):m, 1:p)
+        ΔQ̃ = mul!(ΔQ̃, L₂₁' * ΔL₂₁, Q₁, -1, 1)
+        # Adding ΔA₂ contribution
+        ΔA₂ = mul!(ΔA₂, ΔL₂₁, Q₁, 1, 1)
     end
 
     # construct M
     M = zero!(similar(L, (p, p)))
     if !iszerotangent(ΔL)
-        ΔL11 = view(ΔL, 1:p, 1:p)
-        M = mul!(M, L11', ΔL11, 1, 1)
+        ΔL₁₁ = LowerTriangular(view(ΔL, 1:p, 1:p))
+        M = mul!(M, L₁₁', ΔL₁₁, 1, 1)
     end
-    M = mul!(M, ΔQ̃, Q1', -1, 1)
+    M = mul!(M, ΔQ̃, Q₁', -1, 1)
     view(M, uppertriangularind(M)) .= conj.(view(M, lowertriangularind(M)))
     if eltype(M) <: Complex
         Md = diagview(M)
         Md .= real.(Md)
     end
-    ldiv!(LowerTriangular(L11)', M)
-    ldiv!(LowerTriangular(L11)', ΔQ̃)
-    ΔA1 = mul!(ΔA1, M, Q1, +1, 1)
-    ΔA1 .+= ΔQ̃
+    ldiv!(L₁₁', M)
+    ldiv!(L₁₁', ΔQ̃)
+    ΔA₁ = mul!(ΔA₁, M, Q₁, +1, 1)
+    ΔA₁ .+= ΔQ̃
     return ΔA
 end
 

src/pullbacks/qr.jl

@@ -6,40 +6,26 @@ function check_qr_cotangents(
         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)
-            # 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
-            # cost function
-            ΔQ2 = view(ΔQ, :, (p + 1):size(Q, 2))
-            Δ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_R = norm(ΔR22, Inf)
-            Δgauge = max(Δgauge, Δgauge_R)
-        end
-        Δgauge ≤ gauge_atol ||
-            @warn "`qr` cotangents sensitive to gauge choice: (|Δgauge| = $Δgauge)"
+    Δgauge = abs(zero(eltype(Q)))
+    if !iszerotangent(ΔQ)
+        ΔQ₂ = view(ΔQ, :, (p + 1):minmn)
+        ΔQ₃ = ΔQ[:, (minmn + 1):size(Q, 2)] # extra columns in the case of qr_full
+        Δgauge_Q = norm(ΔQ₂, Inf)
+        Q₁ = view(Q, :, 1:p)
+        Q₁ᴴΔQ₃ = Q₁' * ΔQ₃
+        mul!(ΔQ₃, Q₁, Q₁ᴴΔQ₃, -1, 1)
+        Δgauge_Q = max(Δgauge_Q, norm(ΔQ₃, Inf))
+        Δgauge = max(Δgauge, Δgauge_Q)
+    end
+    if !iszerotangent(ΔR)
+        ΔR22 = view(ΔR, (p + 1):minmn, (p + 1):size(R, 2))
+        Δgauge_R = norm(view(ΔR22, uppertriangularind(ΔR22)), Inf)
+        Δgauge_R = max(Δgauge_R, norm(view(ΔR22, diagind(ΔR22)), Inf))
+        Δgauge = max(Δgauge, Δgauge_R)
     end
-    return
-end
-
-function check_qr_full_cotangents(Q1, ΔQ2, Q1dΔQ2; 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 `qr_full`), there is gauge-invariant information in the
-    # projection of ΔQ2 onto the column space of Q1, by virtue of Q being a unitary
-    # matrix. As the number of Householder reflections is in fixed in the full rank
-    # case, Q is expected to rotate smoothly (we might even be able to predict) also
-    # how the full Q2 will change, but this we omit for now, and we consider
-    # Q2' * ΔQ2 as a gauge dependent quantity.
-    Δgauge = norm(mul!(copy(ΔQ2), Q1, Q1dΔQ2, -1, 1), Inf)
     Δgauge ≤ gauge_atol ||
-        @warn "`qr_full` cotangents sensitive to gauge choice: (|Δgauge| = $Δgauge)"
-    return
+        @warn "`qr` cotangents sensitive to gauge choice: (|Δgauge| = $Δgauge)"
+    return nothing
 end
 
 """
@@ -69,53 +55,54 @@ function qr_pullback!(
     Q, R = QR
     m = size(Q, 1)
     n = size(R, 2)
+    minmn = min(m, n)
     Rd = diagview(R)
     p = qr_rank(R; rank_atol)
 
     ΔQ, ΔR = ΔQR
 
-    Q1 = view(Q, :, 1:p)
-    R11 = view(R, 1:p, 1:p)
-    ΔA1 = view(ΔA, :, 1:p)
-    ΔA2 = view(ΔA, :, (p + 1):n)
+    Q₁ = view(Q, :, 1:p)
+    R₁₁ = UpperTriangular(view(R, 1:p, 1:p))
+    ΔA₁ = view(ΔA, :, 1:p)
+    ΔA₂ = view(ΔA, :, (p + 1):n)
 
     check_qr_cotangents(Q, R, ΔQ, ΔR, p; gauge_atol)
 
     ΔQ̃ = zero!(similar(Q, (m, p)))
     if !iszerotangent(ΔQ)
-        copy!(ΔQ̃, view(ΔQ, :, 1:p))
-        if p < size(Q, 2)
-            Q2 = view(Q, :, (p + 1):size(Q, 2))
-            ΔQ2 = view(ΔQ, :, (p + 1):size(Q, 2))
-            Q1dΔQ2 = Q1' * ΔQ2
-            check_qr_full_cotangents(Q1, ΔQ2, Q1dΔQ2; gauge_atol)
-            ΔQ̃ = mul!(ΔQ̃, Q2, Q1dΔQ2', -1, 1)
+        ΔQ₁ = view(ΔQ, :, 1:p)
+        copy!(ΔQ̃, ΔQ₁)
+        if minmn < size(Q, 2)
+            ΔQ₃ = view(ΔQ, :, (minmn + 1):size(ΔQ, 2)) # extra columns in the case of qr_full
+            Q₃ = view(Q, :, (minmn + 1):size(Q, 2))
+            Q₁ᴴΔQ₃ = Q₁' * ΔQ₃
+            ΔQ̃ = mul!(ΔQ̃, Q₃, Q₁ᴴΔQ₃', -1, 1)
         end
     end
     if !iszerotangent(ΔR) && n > p
-        R12 = view(R, 1:p, (p + 1):n)
-        ΔR12 = view(ΔR, 1:p, (p + 1):n)
-        ΔQ̃ = mul!(ΔQ̃, Q1, ΔR12 * R12', -1, 1)
-        # Adding ΔA2 contribution
-        ΔA2 = mul!(ΔA2, Q1, ΔR12, 1, 1)
+        R₁₂ = view(R, 1:p, (p + 1):n)
+        ΔR₁₂ = view(ΔR, 1:p, (p + 1):n)
+        ΔQ̃ = mul!(ΔQ̃, Q₁, ΔR₁₂ * R₁₂', -1, 1)
+        # Adding ΔA₂ contribution
+        ΔA₂ = mul!(ΔA₂, Q₁, ΔR₁₂, 1, 1)
     end
 
     # construct M
     M = zero!(similar(R, (p, p)))
     if !iszerotangent(ΔR)
-        ΔR11 = view(ΔR, 1:p, 1:p)
-        M = mul!(M, ΔR11, R11', 1, 1)
+        ΔR₁₁ = UpperTriangular(view(ΔR, 1:p, 1:p))
+        M = mul!(M, ΔR₁₁, R₁₁', 1, 1)
     end
-    M = mul!(M, Q1', ΔQ̃, -1, 1)
+    M = mul!(M, Q₁', ΔQ̃, -1, 1)
     view(M, lowertriangularind(M)) .= conj.(view(M, uppertriangularind(M)))
     if eltype(M) <: Complex
         Md = diagview(M)
         Md .= real.(Md)
     end
-    rdiv!(M, UpperTriangular(R11)')
-    rdiv!(ΔQ̃, UpperTriangular(R11)')
-    ΔA1 = mul!(ΔA1, Q1, M, +1, 1)
-    ΔA1 .+= ΔQ̃
+    rdiv!(M, R₁₁') # R₁₁ is upper triangular
+    rdiv!(ΔQ̃, R₁₁')
+    ΔA₁ = mul!(ΔA₁, Q₁, M, +1, 1)
+    ΔA₁ .+= ΔQ̃
     return ΔA
 end
 

src/pullbacks/svd.jl

@@ -1,4 +1,4 @@
-svd_rank(S, rank_atol) = searchsortedlast(S, rank_atol; rev = true)
+svd_rank(S; rank_atol = default_pullback_rank_atol(S)) = searchsortedlast(S, rank_atol; rev = true)
 
 function check_svd_cotangents(aUΔU, Sr, aVΔV; degeneracy_atol = default_pullback_rank_atol(Sr), gauge_atol = default_pullback_gauge_atol(aUΔU, aVΔV))
     mask = abs.(Sr' .- Sr) .< degeneracy_atol
@@ -43,7 +43,7 @@ function svd_pullback!(
     minmn = min(m, n)
     S = diagview(Smat)
     length(S) == minmn || throw(DimensionMismatch("length of S ($(length(S))) does not matrix minimum dimension of U, Vᴴ ($minmn)"))
-    r = svd_rank(S, rank_atol)
+    r = svd_rank(S; rank_atol)
     Ur = view(U, :, 1:r)
     Vᴴr = view(Vᴴ, 1:r, :)
     Sr = view(S, 1:r)

test/enzyme/eig.jl

@@ -14,7 +14,7 @@ is_buildkite = get(ENV, "BUILDKITE", "false") == "true"
 
 m = 19
 for T in (BLASFloats..., GenericFloats...)
-    TestSuite.seed_rng!(123)
+    TestSuite.seed_rng!(1234)
     if !is_buildkite
         TestSuite.test_enzyme_eig(T, (m, m); atol = m * m * TestSuite.precision(T), rtol = m * m * TestSuite.precision(T))
     end

test/enzyme/eigh.jl

@@ -14,7 +14,7 @@ is_buildkite = get(ENV, "BUILDKITE", "false") == "true"
 
 m = 19
 for T in (BLASFloats..., GenericFloats...)
-    TestSuite.seed_rng!(123)
+    TestSuite.seed_rng!(1234)
     if !is_buildkite
         TestSuite.test_enzyme_eigh(T, (m, m); atol = m * m * TestSuite.precision(T), rtol = m * m * TestSuite.precision(T))
     end

test/enzyme/lq.jl

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

test/enzyme/orthnull.jl

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

test/enzyme/polar.jl

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