Fix x64 test

majnemer · Jun 4, 2020 · 38d4837 · 38d4837
1 parent 0f00327
commit 38d4837
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/masking_test.py b/tests/masking_test.py
@@ -413,9 +413,9 @@ def fun(x):
       @jit
       def concat(y):
         return lax.concatenate([x, y], 0)
-      return concat(jnp.array([1., 2., 3.]))
+      return concat(jnp.array([1., 2., 3.], dtype='float32'))
 
-    self.check(fun, ['n'], '(n+3,)', {'n': 2}, [(3,)], ['float_'],
+    self.check(fun, ['n'], '(n+3,)', {'n': 2}, [(3,)], ['float32'],
                rand_default(self.rng()))
 
   @parameterized.named_parameters({