test: fix enzyme frontend tests

SciML · Oct 4, 2024 · b48eb92 · b48eb92
1 parent 087d59c
commit b48eb92
Showing 1 changed file with 24 additions and 25 deletions.
diff --git a/test/alternative_ad_frontend.jl b/test/alternative_ad_frontend.jl
@@ -48,14 +48,13 @@ dup = Zygote.gradient(senseloss(InterpolatingAdjoint()), u0p)[1]
 
 @test ForwardDiff.gradient(senseloss(InterpolatingAdjoint()), u0p) ≈ dup
 
-@test Enzyme.gradient(Reverse, senseloss(InterpolatingAdjoint()), u0p) ≈ dup
-@test_throws SciMLSensitivity.EnzymeTrackedRealError Enzyme.gradient(
-    Reverse, senseloss(ReverseDiffAdjoint()), u0p)≈dup
-@test_throws SciMLSensitivity.EnzymeTrackedRealError Enzyme.gradient(
-    Reverse, senseloss(TrackerAdjoint()), u0p)≈dup
-@test Enzyme.gradient(Reverse, senseloss(ForwardDiffSensitivity()), u0p) ≈ dup
-@test_broken Enzyme.gradient(Reverse, senseloss(ForwardSensitivity()),
-    u0p) ≈ dup # broken because ForwardSensitivity not compatible with perturbing u0
+@test only(Enzyme.gradient(Reverse, senseloss(InterpolatingAdjoint()), u0p)) ≈ dup
+@test_throws SciMLSensitivity.EnzymeTrackedRealError only(Enzyme.gradient(
+    Reverse, senseloss(ReverseDiffAdjoint()), u0p))≈dup
+@test_throws SciMLSensitivity.EnzymeTrackedRealError only(Enzyme.gradient(
+    Reverse, senseloss(TrackerAdjoint()), u0p))≈dup
+@test only(Enzyme.gradient(Reverse, senseloss(ForwardDiffSensitivity()), u0p)) ≈ dup
+@test_broken only(Enzyme.gradient(Reverse, senseloss(ForwardSensitivity()), u0p)) ≈ dup # broken because ForwardSensitivity not compatible with perturbing u0
 
 struct senseloss2{T}
     sense::T
@@ -85,11 +84,11 @@ dup = Zygote.gradient(senseloss2(InterpolatingAdjoint()), u0p)[1]
 
 @test ForwardDiff.gradient(senseloss2(InterpolatingAdjoint()), u0p) ≈ dup
 
-@test_broken Enzyme.gradient(Reverse, senseloss2(InterpolatingAdjoint()), u0p) ≈ dup
-@test_broken Enzyme.gradient(Reverse, senseloss2(ReverseDiffAdjoint()), u0p) ≈ dup
-@test_broken Enzyme.gradient(Reverse, senseloss2(TrackerAdjoint()), u0p) ≈ dup
-@test_broken Enzyme.gradient(Reverse, senseloss2(ForwardDiffSensitivity()), u0p) ≈ dup
-@test_broken Enzyme.gradient(Reverse, senseloss2(ForwardSensitivity()), u0p) ≈ dup # broken because ForwardSensitivity not compatible with perturbing u0
+@test_broken only(Enzyme.gradient(Reverse, senseloss2(InterpolatingAdjoint()), u0p)) ≈ dup
+@test_broken only(Enzyme.gradient(Reverse, senseloss2(ReverseDiffAdjoint()), u0p)) ≈ dup
+@test_broken only(Enzyme.gradient(Reverse, senseloss2(TrackerAdjoint()), u0p)) ≈ dup
+@test_broken only(Enzyme.gradient(Reverse, senseloss2(ForwardDiffSensitivity()), u0p)) ≈ dup
+@test_broken only(Enzyme.gradient(Reverse, senseloss2(ForwardSensitivity()), u0p)) ≈ dup # broken because ForwardSensitivity not compatible with perturbing u0
 
 struct senseloss3{T}
     sense::T
@@ -117,11 +116,11 @@ dup = Zygote.gradient(senseloss3(InterpolatingAdjoint()), u0p)[1]
 
 @test ForwardDiff.gradient(senseloss3(InterpolatingAdjoint()), u0p) ≈ dup
 
-@test_broken Enzyme.gradient(Reverse, senseloss3(InterpolatingAdjoint()), u0p) ≈ dup
-@test_broken Enzyme.gradient(Reverse, senseloss3(ReverseDiffAdjoint()), u0p) ≈ dup
-@test_broken Enzyme.gradient(Reverse, senseloss3(TrackerAdjoint()), u0p) ≈ dup
-@test_broken Enzyme.gradient(Reverse, senseloss3(ForwardDiffSensitivity()), u0p) ≈ dup
-@test_broken Enzyme.gradient(Reverse, senseloss3(ForwardSensitivity()), u0p) ≈ dup
+@test_broken only(Enzyme.gradient(Reverse, senseloss3(InterpolatingAdjoint()), u0p)) ≈ dup
+@test_broken only(Enzyme.gradient(Reverse, senseloss3(ReverseDiffAdjoint()), u0p)) ≈ dup
+@test_broken only(Enzyme.gradient(Reverse, senseloss3(TrackerAdjoint()), u0p)) ≈ dup
+@test_broken only(Enzyme.gradient(Reverse, senseloss3(ForwardDiffSensitivity()), u0p)) ≈ dup
+@test_broken only(Enzyme.gradient(Reverse, senseloss3(ForwardSensitivity()), u0p)) ≈ dup
 
 struct senseloss4{T}
     sense::T
@@ -149,13 +148,13 @@ dup = Zygote.gradient(senseloss4(InterpolatingAdjoint()), u0p)[1]
 
 @test ForwardDiff.gradient(senseloss4(InterpolatingAdjoint()), u0p) ≈ dup
 
-@test Enzyme.gradient(Reverse, senseloss4(InterpolatingAdjoint()), u0p) ≈ dup
-@test_throws SciMLSensitivity.EnzymeTrackedRealError Enzyme.gradient(
-    Reverse, senseloss4(ReverseDiffAdjoint()), u0p)≈dup
-@test_throws SciMLSensitivity.EnzymeTrackedRealError Enzyme.gradient(
-    Reverse, senseloss4(TrackerAdjoint()), u0p)≈dup
-@test Enzyme.gradient(Reverse, senseloss4(ForwardDiffSensitivity()), u0p) ≈ dup
-@test_broken Enzyme.gradient(Reverse, senseloss4(ForwardSensitivity()), u0p) ≈ dup
+@test only(Enzyme.gradient(Reverse, senseloss4(InterpolatingAdjoint()), u0p)) ≈ dup
+@test_throws SciMLSensitivity.EnzymeTrackedRealError only(Enzyme.gradient(
+    Reverse, senseloss4(ReverseDiffAdjoint()), u0p))≈dup
+@test_throws SciMLSensitivity.EnzymeTrackedRealError only(Enzyme.gradient(
+    Reverse, senseloss4(TrackerAdjoint()), u0p))≈dup
+@test only(Enzyme.gradient(Reverse, senseloss4(ForwardDiffSensitivity()), u0p)) ≈ dup
+@test_broken only(Enzyme.gradient(Reverse, senseloss4(ForwardSensitivity()), u0p)) ≈ dup # broken because ForwardSensitivity not compatible with perturbing u0
 
 solvealg_test = Tsit5()
 sensealg_test = InterpolatingAdjoint()