Fix svd_trunc

Author: kshyatt

ext/MatrixAlgebraKitAMDGPUExt/yarocsolver.jl

@@ -164,7 +164,7 @@ for (fname, elty, relty) in
 
             AMDGPU.unsafe_free!(dev_residual) 
             AMDGPU.unsafe_free!(dev_n_sweeps) 
-            return U, S, Vᴴ
+            return (S, U, Vᴴ)
         end
     end
 end

ext/MatrixAlgebraKitCUDAExt/yacusolver.jl

@@ -242,7 +242,7 @@ for (bname, fname, elty, relty) in
             if jobz == 'V'
                 adjoint!(Vᴴ, Ṽ)
             end
-            return U, S, Vᴴ
+            return S, U, Vᴴ
         end
     end
 end

src/common/gauge.jl

@@ -1,7 +1,7 @@
 function gaugefix!(V::AbstractMatrix)
     for j in axes(V, 2)
         v = view(V, :, j)
-        s = conj(sign(argmax(abs, v)))
+        s = conj(sign(_argmaxabs(v)))
         @inbounds v .*= s
     end
     return V
@@ -12,16 +12,16 @@ function gaugefix!(::Val{:full}, U, S, Vᴴ, m::Int, n::Int)
         if j <= min(m, n)
             u = view(U, :, j)
             v = view(Vᴴ, j, :)
-            s = conj(sign(argmax(abs, u)))
+            s = conj(sign(_argmaxabs(u)))
             u .*= s
             v .*= conj(s)
         elseif j <= m
             u = view(U, :, j)
-            s = conj(sign(argmax(abs, u)))
+            s = conj(sign(_argmaxabs(u)))
             u .*= s
         else
             v = view(Vᴴ, j, :)
-            s = conj(sign(argmax(abs, v)))
+            s = conj(sign(_argmaxabs(v)))
             v .*= s
         end
     end
@@ -32,7 +32,18 @@ function gaugefix!(::Val{:compact}, U, S, Vᴴ, m::Int, n::Int)
     for j in 1:size(U, 2)
         u = view(U, :, j)
         v = view(Vᴴ, j, :)
-        s = conj(sign(argmax(abs, u)))
+        s = conj(sign(_argmaxabs(u)))
+        u .*= s
+        v .*= conj(s)
+    end
+    return (U, S, Vᴴ)
+end
+
+function gaugefix!(::Val{:trunc}, U, S, Vᴴ, m::Int, n::Int)
+    for j in 1:min(m, n)
+        u = view(U, :, j)
+        v = view(Vᴴ, j, :)
+        s = conj(sign(_argmaxabs(u)))
         u .*= s
         v .*= conj(s)
     end

src/implementations/svd.jl

@@ -175,21 +175,21 @@ const GPU_SVDPolar = Union{CUSOLVER_SVDPolar}
 const GPU_Jacobi = Union{CUSOLVER_Jacobi, ROCSOLVER_Jacobi}
 const GPU_Randomized = Union{CUSOLVER_Randomized}
 
-function check_input(::typeof(svd_trunc!), A::AbstractMatrix, USVᴴ, alg::TruncatedAlgorithm{CUSOLVER_Randomized})
+function check_input(::typeof(svd_trunc!), A::AbstractMatrix, USVᴴ, alg::CUSOLVER_Randomized)
     m, n = size(A)
     minmn = min(m, n)
     U, S, Vᴴ = USVᴴ
     @assert U isa AbstractMatrix && S isa Diagonal && Vᴴ isa AbstractMatrix
     @check_size(U, (m, m))
     @check_scalar(U, A)
-    @check_size(S, (minmn,minmn))
+    @check_size(S, (minmn, minmn))
     @check_scalar(S, A, real)
     @check_size(Vᴴ, (n, n))
     @check_scalar(Vᴴ, A)
     return nothing
 end
 
-function initialize_output(::typeof(svd_trunc!), A::AbstractMatrix, alg::TruncatedAlgorithm{CUSOLVER_Randomized})
+function initialize_output(::typeof(svd_trunc!), A::AbstractMatrix, alg::TruncatedAlgorithm{<:CUSOLVER_Randomized})
     m, n = size(A)
     minmn = min(m, n)
     U = similar(A, (m, m))
@@ -232,12 +232,12 @@ function MatrixAlgebraKit.svd_full!(A::AbstractMatrix, USVᴴ, alg::GPU_SVDAlgor
 end
 
 function svd_trunc!(A::AbstractMatrix, USVᴴ, alg::TruncatedAlgorithm{<:GPU_Randomized})
-    check_input(svd_trunc!, A, USVᴴ, alg)
+    check_input(svd_trunc!, A, USVᴴ, alg.alg)
     U, S, Vᴴ = USVᴴ
-    _gpu_Xgesvdr!(A, S.diag, U, Vᴴ; alg.kwargs...)
+    _gpu_Xgesvdr!(A, S.diag, U, Vᴴ; alg.alg.kwargs...)
     # TODO: make this controllable using a `gaugefix` keyword argument
-    gaugefix!(Val(:compact), U, S, Vᴴ, m, n)
-    return truncate!(svd_trunc!, USVᴴ′, alg.trunc)
+    gaugefix!(Val(:trunc), U, S, Vᴴ, size(A)...)
+    return truncate!(svd_trunc!, USVᴴ, alg.trunc)
 end
 
 function MatrixAlgebraKit.svd_compact!(A::AbstractMatrix, USVᴴ, alg::GPU_SVDAlgorithm)
@@ -255,7 +255,7 @@ function MatrixAlgebraKit.svd_compact!(A::AbstractMatrix, USVᴴ, alg::GPU_SVDAl
         throw(ArgumentError("Unsupported SVD algorithm"))
     end
     # TODO: make this controllable using a `gaugefix` keyword argument
-    gaugefix!(Val(:compact), U, S, Vᴴ, m, n)
+    gaugefix!(Val(:compact), U, S, Vᴴ, size(A)...) 
     return USVᴴ
 end
 _argmaxabs(x) = reduce(_largest, x; init=zero(eltype(x)))