Add projections tests for Enzyme

Author: Katharine Hyatt

test/enzyme/projections.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...)
+    TestSuite.seed_rng!(123)
+    atol = rtol = m * m * TestSuite.precision(T)
+    if !is_buildkite
+        TestSuite.test_enzyme_projections(T, (m, m); atol, rtol)
+        TestSuite.test_enzyme_projections(Diagonal{T, Vector{T}}, (m, m); atol, rtol)
+    end
+end

test/testsuite/TestSuite.jl

@@ -130,5 +130,6 @@ include("enzyme/lq.jl")
 include("enzyme/svd.jl")
 include("enzyme/polar.jl")
 include("enzyme/orthnull.jl")
+include("enzyme/projections.jl")
 
 end

test/testsuite/enzyme/projections.jl

@@ -0,0 +1,53 @@
+"""
+    test_enzyme_projections(T, sz; kwargs...)
+
+Run all Enzyme AD tests for hermitian and anti-hermitian projections of element type `T`
+and size `sz`.
+"""
+function test_enzyme_projections(T::Type, sz; kwargs...)
+    summary_str = testargs_summary(T, sz)
+    return @testset "Enzyme projection $summary_str" begin
+        test_enzyme_project_hermitian(T, sz; kwargs...)
+        test_enzyme_project_antihermitian(T, sz; kwargs...)
+    end
+end
+
+"""
+    test_enzyme_project_hermitian(T, sz; rng, atol, rtol)
+
+Test the Enzyme reverse-mode AD rule for `project_hermitian` and its in-place variant.
+"""
+function test_enzyme_project_hermitian(
+        T, sz;
+        rng = Random.default_rng(), atol::Real = 0, rtol::Real = precision(T),
+        fdm = enzyme_fdm(T)
+    )
+    return @testset "project_hermitian" begin
+        A = instantiate_matrix(T, sz)
+        B = instantiate_matrix(T, sz)
+        alg = MatrixAlgebraKit.select_algorithm(project_hermitian, A)
+        test_reverse(project_hermitian, RT, (A, TA), (alg, Const); atol, rtol, fdm)
+        test_reverse(project_hermitian!, RT, (A, TA), (A, TA), (alg, Const); atol, rtol, fdm)
+        test_reverse(project_hermitian!, RT, (A, TA), (B, TA), (alg, Const); atol, rtol, fdm)
+    end
+end
+
+"""
+    test_enzyme_project_antihermitian(T, sz; rng, atol, rtol)
+
+Test the Enzyme reverse-mode AD rule for `project_antihermitian` and its in-place variant.
+"""
+function test_enzyme_project_antihermitian(
+        T, sz;
+        rng = Random.default_rng(), atol::Real = 0, rtol::Real = precision(T),
+        fdm = enzyme_fdm(T)
+    )
+    return @testset "project_antihermitian" begin
+        A = instantiate_matrix(T, sz)
+        B = instantiate_matrix(T, sz)
+        alg = MatrixAlgebraKit.select_algorithm(project_hermitian, A)
+        test_reverse(project_antihermitian, RT, (A, TA), (alg, Const); atol, rtol, fdm)
+        test_reverse(project_antihermitian!, RT, (A, TA), (A, TA), (alg, Const); atol, rtol, fdm)
+        test_reverse(project_antihermitian!, RT, (A, TA), (B, TA), (alg, Const); atol, rtol, fdm)
+    end
+end