Reproducer

Author: kshyatt

Project.toml

@@ -61,3 +61,6 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
 test = ["Aqua", "JET", "SafeTestsets", "Test", "TestExtras", "ChainRulesCore", "ChainRulesTestUtils", "StableRNGs", "Zygote", "CUDA", "AMDGPU", "GenericLinearAlgebra", "GenericSchur", "Mooncake", "Enzyme", "EnzymeTestUtils"]
+
+[sources]
+EnzymeTestUtils = {path="/Users/khyatt/.julia/dev/Enzyme/lib/EnzymeTestUtils"}

ext/MatrixAlgebraKitEnzymeExt/MatrixAlgebraKitEnzymeExt.jl

@@ -190,18 +190,22 @@ function EnzymeRules.augmented_primal(
         func::Const{typeof(svd_trunc!)},
         ::Type{RT},
         A::Annotation,
-        USVᴴ::Annotation,
+        USVᴴϵ::Annotation,
         alg::Const{<:MatrixAlgebraKit.AbstractAlgorithm},
     ) where {RT}
     # form cache if needed
     cache_A = copy(A.val)
-    svd_compact!(A.val, USVᴴ.val, alg.val.alg)
-    cache_USVᴴ = copy.(USVᴴ.val)
-    USVᴴ′, ind = MatrixAlgebraKit.truncate(svd_trunc!, USVᴴ.val, alg.val.trunc)
-    ϵ.val      = MatrixAlgebraKit.truncation_error!(diagview(USVᴴ.val[2]), ind)
-    primal     = EnzymeRules.needs_primal(config) ? (USVᴴ′..., ϵ.val) : nothing
-    shadow_USVᴴ = if !isa(A, Const) && !isa(USVᴴ, Const)
-        dU, dS, dVᴴ = USVᴴ.dval
+    USVᴴ = USVᴴϵ.val[1:3]
+    ϵ    = USVᴴϵ.val[end]
+    svd_compact!(A.val, USVᴴ, alg.val.alg)
+    cache_USVᴴ = copy.(USVᴴ)
+    USVᴴ′, ind = MatrixAlgebraKit.truncate(svd_trunc!, USVᴴ, alg.val.trunc)
+    if !isempty(ϵ)
+        ϵ .= MatrixAlgebraKit.truncation_error!(diagview(USVᴴ[2]), ind)
+    end
+    primal     = EnzymeRules.needs_primal(config) ? (USVᴴ′..., ϵ) : nothing
+    shadow_USVᴴ = if !isa(A, Const) && !isa(USVᴴϵ, Const)
+        dU, dS, dVᴴ, dϵ = USVᴴϵ.dval
         # This creates new output shadow matrices, we do this slicing
         # to ensure they have the correct eltype and dimensions.
         # These new shadow matrices are "filled in" with the accumulated
@@ -214,7 +218,7 @@ function EnzymeRules.augmented_primal(
     else
         (nothing, nothing, nothing)
     end
-    shadow = EnzymeRules.needs_shadow(config) ? (shadow_USVᴴ..., ϵ.dval) : nothing
+    shadow = EnzymeRules.needs_shadow(config) ? (shadow_USVᴴ..., USVᴴϵ.dval[end]) : nothing
     return EnzymeRules.AugmentedReturn(primal, shadow, (cache_A, cache_USVᴴ, shadow_USVᴴ, ind))
 end
 function EnzymeRules.reverse(
@@ -223,7 +227,7 @@ function EnzymeRules.reverse(
         dret::Type{RT},
         cache,
         A::Annotation,
-        USVᴴ::Annotation,
+        USVᴴϵ::Annotation,
         alg::Const{<:MatrixAlgebraKit.AbstractAlgorithm},
     ) where {RT}
     cache_A, cache_USVᴴ, shadow_USVᴴ, ind = cache
@@ -233,8 +237,7 @@ function EnzymeRules.reverse(
     if !isa(A, Const) && !isa(USVᴴ, Const)
         svd_pullback!(A.dval, Aval, (U, S, Vᴴ), shadow_USVᴴ, ind)
     end
-    !isa(USVᴴ, Const) && make_zero!(USVᴴ.dval)
-    !isa(ϵ, Const) && make_zero!(ϵ.dval)
+    !isa(USVᴴϵ, Const) && make_zero!(USVᴴϵ.dval)
     return (nothing, nothing, nothing, nothing)
 end
 

test/enzyme.jl

@@ -9,7 +9,8 @@ using LinearAlgebra: UpperTriangular, Diagonal, Hermitian, mul!
 
 is_ci = get(ENV, "CI", "false") == "true"
 
-ETs = is_ci ? (Float64, Float32) : (Float64, Float32, ComplexF32, ComplexF64) # Enzyme/#2631
+#ETs = is_ci ? (Float64, Float32) : (Float64, Float32, ComplexF32, ComplexF64) # Enzyme/#2631
+ETs = (Float64,) # Enzyme/#2631
 include("ad_utils.jl")
 function test_pullbacks_match(rng, f!, f, A, args, Δargs, alg = nothing; ȳ = copy.(Δargs), return_act = Duplicated)
     ΔA = randn(rng, eltype(A), size(A)...)
@@ -180,7 +181,7 @@ end
             test_reverse(eig_vals, RT, (A, TA); fkwargs = (alg = alg,), atol = atol, rtol = rtol, output_tangent = copy(ΔD2.diag))
             test_pullbacks_match(rng, eig_vals!, eig_vals, A, D.diag, ΔD2.diag, alg)
         end
-        @testset "eig_trunc reverse: RT $RT, TA $TA" for RT in (MixedDuplicated,), TA in (Duplicated,)
+        @testset "eig_trunc reverse: RT $RT, TA $TA" for RT in (Duplicated,), TA in (Duplicated,)
             for r in 1:4:m
                 truncalg = TruncatedAlgorithm(alg, truncrank(r; by = abs))
                 ind = MatrixAlgebraKit.findtruncated(diagview(D), truncalg.trunc)
@@ -298,7 +299,7 @@ end
             test_reverse(copy_eigh_vals, RT, (A, TA); fkwargs = (alg = alg,), atol = atol, rtol = rtol, output_tangent = copy(ΔD2.diag))
             test_pullbacks_match(rng, copy_eigh_vals!, copy_eigh_vals, A, D.diag, ΔD2.diag, alg)
         end
-        @testset "eigh_trunc reverse: RT $RT, TA $TA" for RT in (MixedDuplicated,), TA in (Duplicated,)
+        @testset "eigh_trunc reverse: RT $RT, TA $TA" for RT in (Duplicated,), TA in (Duplicated,)
             for r in 1:4:m
                 Ddiag = diagview(D)
                 truncalg = TruncatedAlgorithm(alg, truncrank(r; by = abs))
@@ -334,13 +335,13 @@ end
 @timedtestset "SVD AD Rules with eltype $T" for T in ETs
     rng = StableRNG(12345)
     m = 19
-    @testset "size ($m, $n)" for n in (17, m, 23)
+    @testset "size ($m, $n)" for n in (17,)# m, 23)
         atol = rtol = m * n * precision(T)
         A = randn(rng, T, m, n)
         minmn = min(m, n)
         @testset for alg in (
                 LAPACK_QRIteration(),
-                LAPACK_DivideAndConquer(),
+                #LAPACK_DivideAndConquer(),
             )
             #=@testset "reverse: RT $RT, TA $TA" for RT in (Duplicated,), TA in (Duplicated,)
                 @testset "svd_compact" begin
@@ -378,8 +379,9 @@ end
                     test_pullbacks_match(rng, svd_vals!, svd_vals, A, S, ΔS, alg)
                 end
             end=#
-            @testset "svd_trunc reverse: RT $RT, TA $TA" for RT in (MixedDuplicated,), TA in (Duplicated,)
-                for r in 1:4:minmn
+            fdm = T <: Union{Float32, ComplexF32} ? EnzymeTestUtils.FiniteDifferences.central_fdm(5, 1, max_range = 1.0e-2) : EnzymeTestUtils.FiniteDifferences.central_fdm(5, 1)
+            @testset "svd_trunc reverse: RT $RT, TA $TA" for RT in (Duplicated,), TA in (Duplicated,)
+                #=for r in 1:4:minmn
                     U, S, Vᴴ = svd_compact(A)
                     ΔU = randn(rng, T, m, minmn)
                     ΔS = randn(rng, real(T), minmn, minmn)
@@ -394,11 +396,9 @@ end
                     ΔStrunc = Diagonal(diagview(ΔS2)[ind])
                     ΔUtrunc = ΔU[:, ind]
                     ΔVᴴtrunc = ΔVᴴ[ind, :]
-                    fdm = T <: Union{Float32, ComplexF32} ? EnzymeTestUtils.FiniteDifferences.central_fdm(5, 1, max_range = 1.0e-2) : EnzymeTestUtils.FiniteDifferences.central_fdm(5, 1)
-                    # broken due to Enzyme -- copying in gaugefix????
-                    test_reverse(svd_trunc, RT, (A, TA); fkwargs = (alg = truncalg,), atol = atol, rtol = rtol, output_tangent = (ΔUtrunc, ΔStrunc, ΔVᴴtrunc, zero(real(T))), fdm = fdm)
-                    test_pullbacks_match(rng, svd_trunc!, svd_trunc, A, (U, S, Vᴴ), (ΔU, ΔS2, ΔVᴴ), truncalg, ȳ=(ΔUtrunc, ΔStrunc, ΔVᴴtrunc, zero(real(T))), return_act=RT)
-                end
+                    test_reverse(svd_trunc!, RT, (copy(A), TA), ((copy(U), copy(S), copy(Vᴴ), [zero(real(T))]), TA), (truncalg, Const); atol = atol, rtol = rtol, output_tangent = (copy(ΔUtrunc), copy(ΔStrunc), copy(ΔVᴴtrunc), [zero(real(T))]), fdm = fdm)
+                    #test_pullbacks_match(rng, svd_trunc!, svd_trunc, A, (U, S, Vᴴ), (ΔU, ΔS2, ΔVᴴ), truncalg, ȳ=(ΔUtrunc, ΔStrunc, ΔVᴴtrunc, zero(real(T))), return_act=RT)
+                end=#
                 U, S, Vᴴ = svd_compact(A)
                 ΔU = randn(rng, T, m, minmn)
                 ΔS = randn(rng, real(T), minmn, minmn)
@@ -413,10 +413,8 @@ end
                 ΔStrunc = Diagonal(diagview(ΔS2)[ind])
                 ΔUtrunc = ΔU[:, ind]
                 ΔVᴴtrunc = ΔVᴴ[ind, :]
-                fdm = T <: Union{Float32, ComplexF32} ? EnzymeTestUtils.FiniteDifferences.central_fdm(5, 1, max_range = 1.0e-2) : EnzymeTestUtils.FiniteDifferences.central_fdm(5, 1)
-                # broken due to Enzyme
-                test_reverse(svd_trunc, RT, (A, TA); fkwargs = (alg = truncalg,), atol = atol, rtol = rtol, output_tangent = (ΔUtrunc, ΔStrunc, ΔVᴴtrunc, zero(real(T))), fdm = fdm)
-                test_pullbacks_match(rng, svd_trunc!, svd_trunc, A, (U, S, Vᴴ), (ΔU, ΔS2, ΔVᴴ), truncalg, ȳ=(ΔUtrunc, ΔStrunc, ΔVᴴtrunc, zero(real(T))), return_act=RT)
+                test_reverse(svd_trunc!, RT, (A, TA), ((copy(U), copy(S), copy(Vᴴ), [zero(real(T))]), TA), (truncalg, Const); atol = atol, rtol = rtol, output_tangent = (ΔUtrunc, ΔStrunc, ΔVᴴtrunc, [zero(real(T))]), fdm = fdm)
+                #test_pullbacks_match(rng, svd_trunc!, svd_trunc, A, (U, S, Vᴴ), (ΔU, ΔS2, ΔVᴴ), truncalg, ȳ=(ΔUtrunc, ΔStrunc, ΔVᴴtrunc, zero(real(T))), return_act=RT)
             end
         end
     end