Repository: google/jax

Issue:
Bug in np.nan_to_num in 64-bit mode when x is float32.
Here is a repro:

`
np.nan_to_num(np.array([1, 2, 3], dtype=np.float32))
`

which gives the stack trace:

google3/third_party/py/jax/numpy/lax_numpy.py in nan_to_num(***failed resolving arguments***)
    893     raise ValueError("nan_to_num is not well defined for complex types")
    894   info = finfo(xla_bridge.canonicalize_dtype(_dtype(x)))
--> 895   x = where(isnan(x), _constant_like(x, 0), x)
    896   x = where(isposinf(x), _constant_like(x, info.max), x)
    897   x = where(isneginf(x), _constant_like(x, info.min), x)

google3/third_party/py/jax/numpy/lax_numpy.py in isnan(x)
    885 def isnan(x):
    886   return lax.bitwise_and(lax.bitwise_not(isfinite(x)),
--> 887                          lax.bitwise_not(isinf(x)))
    888 
    889 @_wraps(onp.nan_to_num)

google3/third_party/py/jax/numpy/lax_numpy.py in isinf(x)
    863   dtype = _dtype(x)
    864   if issubdtype(dtype, floating):
--> 865     return lax.eq(lax.abs(x), inf)
    866   elif issubdtype(dtype, complexfloating):
    867     return lax.bitwise_or(lax.eq(lax.abs(real(x)), inf),

google3/third_party/py/jax/lax/lax.py in eq(x, y)
    270 def eq(x, y):
    271   r"""Elementwise equals: :math:`x = y`."""
--> 272   return eq_p.bind(x, y)
    273 
    274 def ne(x, y):

google3/third_party/py/jax/core.py in bind(self, *args, **kwargs)
     73     top_trace = find_top_trace(args)
     74     if top_trace is None:
---> 75       return self.impl(*args, **kwargs)
     76 
     77     tracers = map(top_trace.full_raise, args)

google3/third_party/py/jax/interpreters/xla.py in apply_primitive(prim, *args, **kwargs)
     49 def apply_primitive(prim, *args, **kwargs):
     50   abstract_args = map(abstractify, args)
---> 51   compiled_fun = xla_primitive_callable(prim, *abstract_args, **kwargs)
     52   return compiled_fun(*args)
     53 

google3/third_party/py/jax/util.py in memoized_fun(*args, **kwargs)
    172       cache.popitem(last=False)
    173 
--> 174     ans = cache[key] = fun(*args, **kwargs)
    175     return ans
    176   return memoized_fun

google3/third_party/py/jax/interpreters/xla.py in xla_primitive_callable(prim, *abstract_args, **kwargs)
     59   handle_result = result_handler(result_shape)
     60   compiled = built_c.Compile(shapes, xb.get_compile_options(),
---> 61                              backend=xb.get_backend())
     62   return partial(execute_compiled_primitive, compiled, handle_result)
     63 

google3/third_party/tensorflow/compiler/xla/python/xla_client.py in Compile(self, argument_shapes, compile_options, backend)
    695     if argument_shapes:
    696       compile_options.argument_layouts = argument_shapes
--> 697     c = backend.compile(self.computation, compile_options)
    698     return Executable(c, backend=backend)
    699 

google3/third_party/tensorflow/compiler/xla/python/xla_client.py in compile(self, c_computation, compile_options)
    147     options.debug_options.xla_cpu_fast_math_honor_nans = True
    148     return _xla.LocalExecutable.Compile(c_computation, argument_layouts,
--> 149                                         options, self.client)
    150 
    151   def delete_executable(self, executable):

RuntimeError: Internal: Seen floating point types of different precisions in %compare.4 = pred[3]{0} compare(f32[3]{0} %parameter.1, f64[3]{0} %broadcast.3), direction=EQ, but mixed precision is disallowed.



Discussion:

Resolving patch:
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -850,12 +850,6 @@ def fix(x, out=None):
   zero = lax._const(x, 0)
   return where(lax.ge(x, zero), lax.floor(x), lax.ceil(x))
 
-
-# Caution: If fast math mode is enabled, the semantics of inf and nan are not
-# preserved by XLA/LLVM, and the behavior of inf/nan values is unpredictable.
-# To disable fast math mode on CPU, set the environment variable
-# XLA_FLAGS=--xla_cpu_enable_fast_math=false.
-
 @_wraps(onp.isfinite)
 def isfinite(x):
   dtype = _dtype(x)
@@ -870,17 +864,19 @@ def isfinite(x):
 def isinf(x):
   dtype = _dtype(x)
   if issubdtype(dtype, floating):
-    return lax.eq(lax.abs(x), inf)
+    return lax.eq(lax.abs(x), _constant_like(x, inf))
   elif issubdtype(dtype, complexfloating):
-    return lax.bitwise_or(lax.eq(lax.abs(real(x)), inf),
-                          lax.eq(lax.abs(imag(x)), inf))
+    re = lax.real(x)
+    im = lax.imag(x)
+    return lax.bitwise_or(lax.eq(lax.abs(re), _constant_like(re, inf)),
+                          lax.eq(lax.abs(im), _constant_like(im, inf)))
   else:
     return full_like(x, False, dtype=bool_)
 
 def _isposneginf(infinity, x):
   dtype = _dtype(x)
   if issubdtype(dtype, floating):
-    return lax.eq(x, infinity)
+    return lax.eq(x, _constant_like(x, infinity))
   elif issubdtype(dtype, complexfloating):
     raise ValueError("isposinf/isneginf are not well defined for complex types")
   else:
@@ -897,9 +893,10 @@ def isnan(x):
 @_wraps(onp.nan_to_num)
 def nan_to_num(x, copy=True):
   del copy
-  if iscomplexobj(x):
-    raise ValueError("nan_to_num is not well defined for complex types")
-  info = finfo(xla_bridge.canonicalize_dtype(_dtype(x)))
+  dtype = _dtype(x)
+  if issubdtype(dtype, complexfloating):
+    return lax.complex(nan_to_num(lax.real(x)), nan_to_num(lax.imag(x)))
+  info = finfo(xla_bridge.canonicalize_dtype(dtype))
   x = where(isnan(x), _constant_like(x, 0), x)
   x = where(isposinf(x), _constant_like(x, info.max), x)
   x = where(isneginf(x), _constant_like(x, info.min), x)


Test patch:
diff --git a/jax/test_util.py b/jax/test_util.py
--- a/jax/test_util.py
+++ b/jax/test_util.py
@@ -322,7 +322,6 @@ def post(x):
   return partial(_rand_dtype, randn, scale=100., post=post)
 
 
-# TODO(mattjj): doesn't handle complex types
 def rand_some_inf():
   """Return a random sampler that produces infinities in floating types."""
   rng = npr.RandomState(1)
@@ -334,6 +333,10 @@ def rand(shape, dtype):
       # only float types have inf
       return base_rand(shape, dtype)
 
+    if onp.issubdtype(dtype, onp.complexfloating):
+      base_dtype = onp.real(onp.array(0, dtype=dtype)).dtype
+      return rand(shape, base_dtype) + 1j * rand(shape, base_dtype)
+
     dims = _dims_of_shape(shape)
     posinf_flips = rng.rand(*dims) < 0.1
     neginf_flips = rng.rand(*dims) < 0.1
@@ -346,6 +349,34 @@ def rand(shape, dtype):
 
   return rand
 
+def rand_some_inf_and_nan():
+  """Return a random sampler that produces infinities in floating types."""
+  rng = npr.RandomState(1)
+  base_rand = rand_default()
+
+  def rand(shape, dtype):
+    """The random sampler function."""
+    if not onp.issubdtype(dtype, onp.floating):
+      # only float types have inf
+      return base_rand(shape, dtype)
+
+    if onp.issubdtype(dtype, onp.complexfloating):
+      base_dtype = onp.real(onp.array(0, dtype=dtype)).dtype
+      return rand(shape, base_dtype) + 1j * rand(shape, base_dtype)
+
+    dims = _dims_of_shape(shape)
+    posinf_flips = rng.rand(*dims) < 0.1
+    neginf_flips = rng.rand(*dims) < 0.1
+    nan_flips = rng.rand(*dims) < 0.1
+
+    vals = base_rand(shape, dtype)
+    vals = onp.where(posinf_flips, onp.inf, vals)
+    vals = onp.where(neginf_flips, -onp.inf, vals)
+    vals = onp.where(nan_flips, onp.nan, vals)
+
+    return _cast_to_shape(onp.asarray(vals, dtype=dtype), shape, dtype)
+
+  return rand
 
 # TODO(mattjj): doesn't handle complex types
 def rand_some_zero():
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -33,6 +33,7 @@
 from jax import lax
 from jax import numpy as lnp
 from jax import test_util as jtu
+from jax.lib import xla_bridge
 
 from jax.config import config
 config.parse_flags_with_absl()
@@ -82,7 +83,6 @@ def op_record(name, nargs, dtypes, shapes, rng, diff_modes, test_name=None,
     op_record("floor", 1, float_dtypes, all_shapes, jtu.rand_default(), []),
     op_record("greater", 2, number_dtypes, all_shapes, jtu.rand_some_equal(), []),
     op_record("greater_equal", 2, number_dtypes, all_shapes, jtu.rand_some_equal(), []),
-    op_record("isfinite", 1, number_dtypes, all_shapes, jtu.rand_some_inf(), []),
     op_record("less", 2, number_dtypes, all_shapes, jtu.rand_some_equal(), []),
     op_record("less_equal", 2, number_dtypes, all_shapes, jtu.rand_some_equal(), []),
     op_record("log", 1, number_dtypes, all_shapes, jtu.rand_positive(), ["rev"]),
@@ -137,6 +137,11 @@ def op_record(name, nargs, dtypes, shapes, rng, diff_modes, test_name=None,
     op_record("outer", 2, number_dtypes, all_shapes, jtu.rand_default(), []),
     op_record("imag", 1, number_dtypes, all_shapes, jtu.rand_some_inf(), []),
     op_record("iscomplex", 1, number_dtypes, all_shapes, jtu.rand_some_inf(), []),
+    op_record("isfinite", 1, inexact_dtypes, all_shapes, jtu.rand_some_inf_and_nan(), []),
+    op_record("isinf", 1, inexact_dtypes, all_shapes, jtu.rand_some_inf_and_nan(), []),
+    op_record("isnan", 1, inexact_dtypes, all_shapes, jtu.rand_some_inf_and_nan(), []),
+    op_record("isneginf", 1, float_dtypes, all_shapes, jtu.rand_some_inf_and_nan(), []),
+    op_record("isposinf", 1, float_dtypes, all_shapes, jtu.rand_some_inf_and_nan(), []),
     op_record("isreal", 1, number_dtypes, all_shapes, jtu.rand_some_inf(), []),
     op_record("isrealobj", 1, number_dtypes, all_shapes, jtu.rand_some_inf(), []),
     op_record("log2", 1, number_dtypes, all_shapes, jtu.rand_positive(), ["rev"]),
@@ -1355,6 +1360,19 @@ def testVander(self, shape, dtype, n, increasing, rng):
     self._CheckAgainstNumpy(onp_fun, lnp_fun, args_maker, check_dtypes=False)
     self._CompileAndCheck(lnp_fun, args_maker, check_dtypes=False)
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+        {"testcase_name": jtu.format_test_name_suffix("nan_to_num", [shape],
+                                                      [dtype]),
+         "rng": jtu.rand_some_inf_and_nan(), "shape": shape, "dtype": dtype}
+        for shape in all_shapes
+        for dtype in inexact_dtypes))
+  def testNanToNum(self, rng, shape, dtype):
+    dtype = onp.dtype(xla_bridge.canonicalize_dtype(dtype)).type
+    args_maker = lambda: [rng(shape, dtype)]
+    self._CheckAgainstNumpy(onp.nan_to_num, lnp.nan_to_num, args_maker,
+                            check_dtypes=True)
+    self._CompileAndCheck(lnp.nan_to_num, args_maker, check_dtypes=True)
+
   def testIssue330(self):
     x = lnp.full((1, 1), lnp.array([1])[0])  # doesn't crash
     self.assertEqual(x[0, 0], 1)