Avoid refactorization and cleanup tests to allow for numerical errors due to summation ordering

Author: sharanry

ext/LinearSolveEnzymeExt.jl

@@ -28,17 +28,22 @@ end
 function EnzymeCore.EnzymeRules.forward(func::Const{typeof(LinearSolve.solve!)}, ::Type{RT}, linsolve::EnzymeCore.Annotation{LP}; kwargs...) where {RT, LP <: LinearSolve.LinearCache}
     @assert !(linsolve isa Const)
 
-    A = deepcopy(linsolve.val.A) #mutates after function is applied
+    linsolve = deepcopy(linsolve) #mutates after function is applied
     res = func.val(linsolve.val; kwargs...)
 
     if RT <: Const
         return res
     end
-    
-    dres = deepcopy(res)
+   
+    b = deepcopy(linsolve.val.b)
+
     db = linsolve.dval.b
     dA = linsolve.dval.A
-    dres.u .= A \ (db - dA * res.u)
+
+    linsolve.val.b = db - dA * res.u
+    dres = func.val(linsolve.val; kwargs...)
+
+    linsolve.val.b = b
 
     if RT <: DuplicatedNoNeed
         return dres

test/enzyme.jl

@@ -177,12 +177,6 @@ function fb(b; alg = LUFactorization())
 end
 fb(b1)
 
-manual_jac = map(onehot(b1)) do db
-    y = A \ b1
-    sum(inv(A) * (db - dA*y))
-end |> collect
-@show manual_jac
-
 fd_jac = FiniteDiff.finite_difference_jacobian(fb, b1) |> vec
 @show fd_jac
 
@@ -192,8 +186,7 @@ en_jac = map(onehot(b1)) do db1
 end |> collect
 @show en_jac
 
-@test_broken en_jac ≈ manual_jac
-@test_broken en_jac ≈ fd_jac
+@test en_jac ≈ fd_jac atol=1e-6
 
 function fA(A; alg = LUFactorization())
     prob = LinearProblem(A, b1)
@@ -204,12 +197,6 @@ function fA(A; alg = LUFactorization())
 end
 fA(A)
 
-manual_jac = map(onehot(A)) do dA
-    y = A \ b1
-    sum(inv(A) * (db1 - dA*y))
-end |> collect
-@show manual_jac
-
 fd_jac = FiniteDiff.finite_difference_jacobian(fA, A) |> vec
 @show fd_jac
 
@@ -219,5 +206,4 @@ en_jac = map(onehot(A)) do dA
 end |> collect
 @show en_jac
 
-@test_broken en_jac ≈ manual_jac
-@test_broken en_jac ≈ fd_jac
+@test en_jac ≈ fd_jac atol=1e-6