From cee007378dc18269fca1f85617cac458e0cc42af Mon Sep 17 00:00:00 2001
From: Katharine Hyatt <khyatt@flatironinstitute.org>
Date: Wed, 15 Oct 2025 18:23:16 -0400
Subject: [PATCH] Add tests for polar decomposition for GPU

---
 test/amd/polar.jl  | 85 ++++++++++++++++++++++++++++++++++++++++++++++
 test/cuda/polar.jl | 85 ++++++++++++++++++++++++++++++++++++++++++++++
 test/runtests.jl   |  6 ++++
 3 files changed, 176 insertions(+)
 create mode 100644 test/amd/polar.jl
 create mode 100644 test/cuda/polar.jl

diff --git a/test/amd/polar.jl b/test/amd/polar.jl
new file mode 100644
index 00000000..066c0f78
--- /dev/null
+++ b/test/amd/polar.jl
@@ -0,0 +1,85 @@
+using MatrixAlgebraKit
+using Test
+using TestExtras
+using StableRNGs
+using LinearAlgebra: LinearAlgebra, I, isposdef, Hermitian
+using MatrixAlgebraKit: PolarViaSVD
+using AMDGPU
+
+@testset "left_polar! for T = $T" for T in (Float32, Float64, ComplexF32, ComplexF64)
+    rng = StableRNG(123)
+    m = 54
+    @testset "size ($m, $n)" for n in (37, m)
+        k = min(m, n)
+        svd_algs = (ROCSOLVER_QRIteration(), ROCSOLVER_Jacobi())
+        @testset "algorithm $svd_alg" for svd_alg in svd_algs
+            n < m && svd_alg isa ROCSOLVER_QRIteration && continue
+            A = ROCArray(randn(rng, T, m, n))
+            alg = PolarViaSVD(svd_alg)
+            W, P = left_polar(A; alg)
+            @test W isa ROCMatrix{T} && size(W) == (m, n)
+            @test P isa ROCMatrix{T} && size(P) == (n, n)
+            @test W * P ≈ A
+            @test isisometric(W)
+            # work around extremely strict Julia criteria for Hermiticity
+            @test ishermitian(P; rtol = MatrixAlgebraKit.defaulttol(P)) && isposdef(Hermitian(P))
+
+            Ac = similar(A)
+            W2, P2 = @constinferred left_polar!(copy!(Ac, A), (W, P), alg)
+            @test W2 === W
+            @test P2 === P
+            @test W * P ≈ A
+            @test isisometric(W)
+            # work around extremely strict Julia criteria for Hermiticity
+            @test ishermitian(P; rtol = MatrixAlgebraKit.defaulttol(P)) && isposdef(Hermitian(P))
+
+            noP = similar(P, (0, 0))
+            W2, P2 = @constinferred left_polar!(copy!(Ac, A), (W, noP), alg)
+            @test P2 === noP
+            @test W2 === W
+            @test isisometric(W)
+            P = W' * A # compute P explicitly to verify W correctness
+            @test ishermitian(P; rtol = MatrixAlgebraKit.defaulttol(P))
+            @test isposdef(Hermitian(project_hermitian!(P)))
+        end
+    end
+end
+
+@testset "right_polar! for T = $T" for T in (Float32, Float64, ComplexF32, ComplexF64)
+    rng = StableRNG(123)
+    n = 54
+    @testset "size ($m, $n)" for m in (37, n)
+        k = min(m, n)
+        svd_algs = (ROCSOLVER_QRIteration(), ROCSOLVER_Jacobi())
+        @testset "algorithm $svd_alg" for svd_alg in svd_algs
+            n > m && svd_alg isa ROCSOLVER_QRIteration && continue
+            A = ROCArray(randn(rng, T, m, n))
+            alg = PolarViaSVD(svd_alg)
+            P, Wᴴ = right_polar(A; alg)
+            @test Wᴴ isa ROCMatrix{T} && size(Wᴴ) == (m, n)
+            @test P isa ROCMatrix{T} && size(P) == (m, m)
+            @test P * Wᴴ ≈ A
+            @test isisometric(Wᴴ; side = :right)
+            # work around extremely strict Julia criteria for Hermiticity
+            @test ishermitian(P; rtol = MatrixAlgebraKit.defaulttol(P)) && isposdef(Hermitian(P))
+
+            Ac = similar(A)
+            P2, Wᴴ2 = @constinferred right_polar!(copy!(Ac, A), (P, Wᴴ), alg)
+            @test P2 === P
+            @test Wᴴ2 === Wᴴ
+            @test P * Wᴴ ≈ A
+            @test isisometric(Wᴴ; side = :right)
+            # work around extremely strict Julia criteria for Hermiticity
+            @test ishermitian(P; rtol = MatrixAlgebraKit.defaulttol(P)) && isposdef(Hermitian(P))
+
+            noP = similar(P, (0, 0))
+            P2, Wᴴ2 = @constinferred right_polar!(copy!(Ac, A), (noP, Wᴴ), alg)
+            @test P2 === noP
+            @test Wᴴ2 === Wᴴ
+            @test isisometric(Wᴴ; side = :right)
+            P = A * Wᴴ' # compute P explicitly to verify W correctness
+            @test ishermitian(P; rtol = MatrixAlgebraKit.defaulttol(P))
+            @test isposdef(Hermitian(project_hermitian!(P)))
+        end
+    end
+end
diff --git a/test/cuda/polar.jl b/test/cuda/polar.jl
new file mode 100644
index 00000000..a7a76f27
--- /dev/null
+++ b/test/cuda/polar.jl
@@ -0,0 +1,85 @@
+using MatrixAlgebraKit
+using Test
+using TestExtras
+using StableRNGs
+using LinearAlgebra: LinearAlgebra, I, isposdef, Hermitian
+using MatrixAlgebraKit: PolarViaSVD
+using CUDA
+
+@testset "left_polar! for T = $T" for T in (Float32, Float64, ComplexF32, ComplexF64)
+    rng = StableRNG(123)
+    m = 54
+    @testset "size ($m, $n)" for n in (37, m)
+        k = min(m, n)
+        svd_algs = (CUSOLVER_QRIteration(), CUSOLVER_Jacobi())
+        @testset "algorithm $svd_alg" for svd_alg in svd_algs
+            n < m && svd_alg isa CUSOLVER_QRIteration && continue
+            A = CuArray(randn(rng, T, m, n))
+            alg = PolarViaSVD(svd_alg)
+            W, P = left_polar(A; alg)
+            @test W isa CuMatrix{T} && size(W) == (m, n)
+            @test P isa CuMatrix{T} && size(P) == (n, n)
+            @test W * P ≈ A
+            @test isisometric(W)
+            # work around extremely strict Julia criteria for Hermiticity
+            @test ishermitian(P; rtol = MatrixAlgebraKit.defaulttol(P)) && isposdef(Hermitian(P))
+
+            Ac = similar(A)
+            W2, P2 = @constinferred left_polar!(copy!(Ac, A), (W, P), alg)
+            @test W2 === W
+            @test P2 === P
+            @test W * P ≈ A
+            @test isisometric(W)
+            # work around extremely strict Julia criteria for Hermiticity
+            @test ishermitian(P; rtol = MatrixAlgebraKit.defaulttol(P)) && isposdef(Hermitian(P))
+
+            noP = similar(P, (0, 0))
+            W2, P2 = @constinferred left_polar!(copy!(Ac, A), (W, noP), alg)
+            @test P2 === noP
+            @test W2 === W
+            @test isisometric(W)
+            P = W' * A # compute P explicitly to verify W correctness
+            @test ishermitian(P; rtol = MatrixAlgebraKit.defaulttol(P))
+            @test isposdef(Hermitian(project_hermitian!(P)))
+        end
+    end
+end
+
+@testset "right_polar! for T = $T" for T in (Float32, Float64, ComplexF32, ComplexF64)
+    rng = StableRNG(123)
+    n = 54
+    @testset "size ($m, $n)" for m in (37, n)
+        k = min(m, n)
+        svd_algs = (CUSOLVER_QRIteration(), CUSOLVER_Jacobi())
+        @testset "algorithm $svd_alg" for svd_alg in svd_algs
+            n > m && svd_alg isa CUSOLVER_QRIteration && continue
+            A = CuArray(randn(rng, T, m, n))
+            alg = PolarViaSVD(svd_alg)
+            P, Wᴴ = right_polar(A; alg)
+            @test Wᴴ isa CuMatrix{T} && size(Wᴴ) == (m, n)
+            @test P isa CuMatrix{T} && size(P) == (m, m)
+            @test P * Wᴴ ≈ A
+            @test isisometric(Wᴴ; side = :right)
+            # work around extremely strict Julia criteria for Hermiticity
+            @test ishermitian(P; rtol = MatrixAlgebraKit.defaulttol(P)) && isposdef(Hermitian(P))
+
+            Ac = similar(A)
+            P2, Wᴴ2 = @constinferred right_polar!(copy!(Ac, A), (P, Wᴴ), alg)
+            @test P2 === P
+            @test Wᴴ2 === Wᴴ
+            @test P * Wᴴ ≈ A
+            @test isisometric(Wᴴ; side = :right)
+            # work around extremely strict Julia criteria for Hermiticity
+            @test ishermitian(P; rtol = MatrixAlgebraKit.defaulttol(P)) && isposdef(Hermitian(P))
+
+            noP = similar(P, (0, 0))
+            P2, Wᴴ2 = @constinferred right_polar!(copy!(Ac, A), (noP, Wᴴ), alg)
+            @test P2 === noP
+            @test Wᴴ2 === Wᴴ
+            @test isisometric(Wᴴ; side = :right)
+            P = A * Wᴴ' # compute P explicitly to verify W correctness
+            @test ishermitian(P; rtol = MatrixAlgebraKit.defaulttol(P))
+            @test isposdef(Hermitian(project_hermitian!(P)))
+        end
+    end
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index 16d5fb5b..c0d72578 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -75,6 +75,9 @@ if CUDA.functional()
     @safetestset "CUDA Hermitian Eigenvalue Decomposition" begin
         include("cuda/eigh.jl")
     end
+    @safetestset "CUDA Polar Decomposition" begin
+        include("cuda/polar.jl")
+    end
 end
 
 using AMDGPU
@@ -94,4 +97,7 @@ if AMDGPU.functional()
     @safetestset "AMDGPU Hermitian Eigenvalue Decomposition" begin
         include("amd/eigh.jl")
     end
+    @safetestset "AMDGPU Polar Decomposition" begin
+        include("amd/polar.jl")
+    end
 end