Fixed thrown error for GPU.

google · Jul 21, 2020 · 48ad8f0 · 48ad8f0
1 parent 4c3fac6
commit 48ad8f0
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/jax/experimental/jax2tf/tests/primitives_test.py b/jax/experimental/jax2tf/tests/primitives_test.py
@@ -137,7 +137,8 @@ def test_sort(self, harness: primitive_harness.Harness):
   def test_qr(self, harness: primitive_harness.Harness):
     # See jax.lib.lapack.geqrf for the list of compatible types
     if not harness.params["dtype"] in [jnp.float32, jnp.float64, jnp.complex64, jnp.complex128]:
-      with self.assertRaisesRegex(NotImplementedError, "Unsupported dtype"):
+      expected_error = ValueError if jtu.device_under_test() == "gpu" else NotImplementedError
+      with self.assertRaisesRegex(expected_error, "Unsupported dtype"):
         harness.dyn_fun(*harness.dyn_args_maker(self.rng()))
     elif harness.params["dtype"] in [jnp.float32, jnp.float64]:
       self.ConvertAndCompare(harness.dyn_fun, *harness.dyn_args_maker(self.rng()),