Add tests for other algs and handle cases of algs currently not supported

Author: sharanry

ext/LinearSolveEnzymeExt.jl

@@ -33,7 +33,9 @@ function EnzymeCore.EnzymeRules.forward(func::Const{typeof(LinearSolve.solve!)},
     if RT <: Const
         return res
     end
-   
+    if linsolve.alg isa LinearSolve.AbstractKrylovSubspaceMethod
+        error("Algorithm $(_linsolve.alg) is currently not supported by Enzyme rules on LinearSolve.jl. Please open an issue on LinearSolve.jl detailing which algorithm is missing the adjoint handling")
+    end
     b = deepcopy(linsolve.val.b)
 
     db = linsolve.dval.b

test/enzyme.jl

@@ -1,6 +1,7 @@
 using Enzyme, ForwardDiff
 using LinearSolve, LinearAlgebra, Test
 using FiniteDiff
+using SafeTestsets
 
 n = 4
 A = rand(n, n);
@@ -164,46 +165,53 @@ Enzyme.autodiff(Reverse, f3, Duplicated(copy(A), dA), Duplicated(copy(b1), db1),
 @test db2 ≈ db22
 =#
 
-
 A = rand(n, n);
 dA = zeros(n, n);
 b1 = rand(n);
-function fb(b; alg = LUFactorization())
-    prob = LinearProblem(A, b)
+for alg in (
+    LUFactorization(), 
+    RFLUFactorization(),
+    # KrylovJL_GMRES(), fails
+    )
+    alg_str = string(alg)
+    @show alg_str
+    function fb(b)
+        prob = LinearProblem(A, b)
 
-    sol1 = solve(prob, alg)
+        sol1 = solve(prob, alg)
 
-    sum(sol1.u)
-end
-fb(b1)
+        sum(sol1.u)
+    end
+    fb(b1)
 
-fd_jac = FiniteDiff.finite_difference_jacobian(fb, b1) |> vec
-@show fd_jac
+    fd_jac = FiniteDiff.finite_difference_jacobian(fb, b1) |> vec
+    @show fd_jac
 
-en_jac = map(onehot(b1)) do db1
-    eres = Enzyme.autodiff(Forward, fb, Duplicated(copy(b1), db1))
-    eres[1]
-end |> collect
-@show en_jac
+    en_jac = map(onehot(b1)) do db1
+        eres = Enzyme.autodiff(Forward, fb, Duplicated(copy(b1), db1))
+        eres[1]
+    end |> collect
+    @show en_jac
 
-@test en_jac ≈ fd_jac rtol=1e-6
+    @test en_jac ≈ fd_jac rtol=1e-6
 
-function fA(A; alg = LUFactorization())
-    prob = LinearProblem(A, b1)
+    function fA(A)
+        prob = LinearProblem(A, b1)
 
-    sol1 = solve(prob, alg)
+        sol1 = solve(prob, alg)
 
-    sum(sol1.u)
-end
-fA(A)
+        sum(sol1.u)
+    end
+    fA(A)
 
-fd_jac = FiniteDiff.finite_difference_jacobian(fA, A) |> vec
-@show fd_jac
+    fd_jac = FiniteDiff.finite_difference_jacobian(fA, A) |> vec
+    @show fd_jac
 
-en_jac = map(onehot(A)) do dA
-    eres = Enzyme.autodiff(Forward, fA, Duplicated(copy(A), dA))
-    eres[1]
-end |> collect
-@show en_jac
+    en_jac = map(onehot(A)) do dA
+        eres = Enzyme.autodiff(Forward, fA, Duplicated(copy(A), dA))
+        eres[1]
+    end |> collect
+    @show en_jac
 
-@test en_jac ≈ fd_jac rtol=1e-6
+    @test en_jac ≈ fd_jac rtol=1e-6
+end