fix one more test harness

jax-ml · May 31, 2020 · 4116a9c · 4116a9c
1 parent 460c1d2
commit 4116a9c
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/tests/lax_test.py b/tests/lax_test.py
@@ -2764,8 +2764,9 @@ def f2(x, y):
     expected = onp.array(0.0)
     self.assertAllClose(ans, expected, check_dtypes=False)
 
-    with self.assertRaises(TypeError):
-      lax.stop_gradient(lambda x: x)
+    with core.skip_checks():
+      with self.assertRaises(TypeError):
+        lax.stop_gradient(lambda x: x)
 
   # TODO(mattjj): make this a more systematic test
   def testRemainder(self):