Future-proof view test against signaling NaNs (jax-ml#3178)

BuddenD · May 21, 2020 · bb2127c · bb2127c
1 parent c459280
commit bb2127c
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/jax/test_util.py b/jax/test_util.py
@@ -463,7 +463,7 @@ def _rand_dtype(rand, shape, dtype, scale=1., post=lambda x: x):
   return _cast_to_shape(np.asarray(post(vals), dtype), shape, dtype)
 
 
-def rand_fullrange(rng, standardize_nans=True):
+def rand_fullrange(rng, standardize_nans=False):
   """Random numbers that span the full range of available bits."""
   def gen(shape, dtype, post=lambda x: x):
     dtype = np.dtype(dtype)

diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
@@ -2229,12 +2229,14 @@ def testView(self, shape, a_dtype, dtype):
     if not FLAGS.jax_enable_x64:
       if jnp.dtype(a_dtype).itemsize == 8 or jnp.dtype(dtype).itemsize == 8:
         self.skipTest("x64 types are disabled by jax_enable_x64")
-    rng = jtu.rand_fullrange(self.rng(), standardize_nans=True)
+    rng = jtu.rand_fullrange(self.rng())
     args_maker = lambda: [rng(shape, a_dtype)]
     np_op = lambda x: np.asarray(x).view(dtype)
     jnp_op = lambda x: jnp.asarray(x).view(dtype)
-    self._CheckAgainstNumpy(jnp_op, np_op, args_maker, check_dtypes=True)
-    self._CompileAndCheck(jnp_op, args_maker, check_dtypes=True)
+    # Above may produce signaling nans; ignore warnings from invalid values.
+    with np.errstate(invalid='ignore'):
+      self._CheckAgainstNumpy(jnp_op, np_op, args_maker, check_dtypes=True)
+      self._CompileAndCheck(jnp_op, args_maker, check_dtypes=True)
 
   def testPathologicalFloats(self):
     args_maker = lambda: [np.array([