Add test_throws check for wrong user gradient assertion

KnutAM · Mar 25, 2022 · f151302 · f151302
1 parent 42ab791
commit f151302
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/test/test_ad.jl b/test/test_ad.jl
@@ -325,8 +325,16 @@ S(C) = S(C, μ, Kb)
             @test gradient(ts_ts_ts, xs) ≈ gradient(ts_ts_ts_ana, xs)
 
         end
-
+
+        # Test that AssertionError is thrown for erroneous user functions
+        test_self(x) = x
+        test_self_wronggradient(x) = test_self(x), one(x) # Only true for scalars, not for tensors
+        @implement_gradient test_self test_self_wronggradient
+        @test gradient(test_self, rand()) ≈ 1.0 # Should work fine
+        @test_throws AssertionError gradient(test_self, rand(Tensor{2,3}))
+        @test_throws AssertionError gradient(test_self, rand(SymmetricTensor{2,3}))
     end
 
 
 end # testsection
+