add a test for Enzyme rule correctness

SciML · Sep 22, 2023 · a08386d · a08386d
1 parent ce7ffc0
commit a08386d
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 1 deletion.
diff --git a/Project.toml b/Project.toml
@@ -81,6 +81,8 @@ julia = "1.6"
 
 [extras]
 BlockDiagonals = "0a1fb500-61f7-11e9-3c65-f5ef3456f9f0"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 HYPRE = "b5ffcf37-a2bd-41ab-a3da-4bd9bc8ad771"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
@@ -98,4 +100,4 @@ SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "IterativeSolvers", "InteractiveUtils", "JET", "KrylovKit", "Pkg", "Random", "SafeTestsets", "MultiFloats", "ForwardDiff", "HYPRE", "MPI", "MKL_jll", "BlockDiagonals"]
+test = ["Test", "IterativeSolvers", "InteractiveUtils", "JET", "KrylovKit", "Pkg", "Random", "SafeTestsets", "MultiFloats", "ForwardDiff", "HYPRE", "MPI", "MKL_jll", "BlockDiagonals", "Enzyme", "FiniteDiff"]
diff --git a/test/enzyme.jl b/test/enzyme.jl
@@ -0,0 +1,30 @@
+using Enzyme, FiniteDiff
+using LinearSolve, LinearAlgebra, Test
+
+n = 4
+A = rand(n, n);
+dA = zeros(n, n);
+b1 = rand(n);
+db1 = zeros(n);
+b2 = rand(n);
+db2 = zeros(n);
+
+function f(A, b1, b2; alg = LUFactorization())
+    prob = LinearProblem(A, b1)
+
+    sol1 = solve(prob, alg)
+
+    s1 = sol1.u
+    norm(s1)
+end
+
+f(A, b1, b2) # Uses BLAS
+
+Enzyme.autodiff(Reverse, f, Duplicated(copy(A), dA), Duplicated(copy(b1), db1), Duplicated(copy(b2), db2))
+
+dA2 = FiniteDiff.finite_difference_gradient(x->f(x,b1, b2), copy(A))
+db12 = FiniteDiff.finite_difference_gradient(x->f(A,x, b2), copy(b1))
+
+@test dA ≈ dA2
+@test db1 ≈ db12
+@test db2 == zeros(4)
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -13,6 +13,7 @@ if GROUP == "All" || GROUP == "Core"
     @time @safetestset "Non-Square Tests" include("nonsquare.jl")
     @time @safetestset "SparseVector b Tests" include("sparse_vector.jl")
     @time @safetestset "Default Alg Tests" include("default_algs.jl")
+    VERSION >= v"1.9" && @time @safetestset "Enzyme Derivative Rules" include("enzyme.jl")
     @time @safetestset "Traits" include("traits.jl")
 end