fix jaxtype test cases

jax-ml · Jun 1, 2020 · e9de752 · e9de752
1 parent ac8145f
commit e9de752
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 10 deletions.
diff --git a/tests/api_test.py b/tests/api_test.py
@@ -1765,11 +1765,6 @@ def f(a, b):
     with self.assertRaisesRegex(ValueError, msg):
       g(jnp.ones((1, 1)), b=1)
 
-  def test_primitive_jaxtype_error(self):
-    with self.assertRaisesRegex(
-        TypeError, "Argument .* of type .* is not a valid JAX type"):
-      lax.add(1, 'hi')
-
 
 class JaxprTest(jtu.JaxTestCase):
 

diff --git a/tests/lax_test.py b/tests/lax_test.py
@@ -1704,9 +1704,16 @@ def testDynamicUpdateSliceTypeErrors(self):
                                        (onp.int32(1), onp.int16(2))))
 
   def test_tie_in_error(self):
-    with self.assertRaisesRegex(
-        TypeError, ".* of type .*tuple.* is not a valid JAX type"):
-      api.make_jaxpr(lambda x: lax.tie_in((x, x), 1))(1.)
+    with core.skipping_checks():
+      with self.assertRaisesRegex(
+          TypeError, ".* of type .*tuple.* is not a valid JAX type"):
+        api.make_jaxpr(lambda x: lax.tie_in((x, x), 1))(1.)
+
+  def test_primitive_jaxtype_error(self):
+    with core.skipping_checks():
+      with self.assertRaisesRegex(
+          TypeError, "Argument .* of type .* is not a valid JAX type"):
+        lax.add(1, 'hi')
 
 class LazyConstantTest(jtu.JaxTestCase):
   def _Check(self, make_const, expected):
@@ -2757,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.skipping_checks():
+      with self.assertRaises(TypeError):
+        lax.stop_gradient(lambda x: x)
 
   # TODO(mattjj): make this a more systematic test
   def testRemainder(self):