Canonicalize result dtype to fix double precision problem in ldexp (j…

…ax-ml#4069)
shoyer · Aug 15, 2020 · 1316562 · 1316562
1 parent 1dbdaac
commit 1316562
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
@@ -757,7 +757,7 @@ def _normalize_float(x):
 @_wraps(np.ldexp)
 @jit
 def ldexp(x1, x2):
-  dtype = _result_dtype(np.ldexp, x1, x2)
+  dtype = dtypes.canonicalize_dtype(_result_dtype(np.ldexp, x1, x2))
   x1, x2 = _promote_shapes("ldexp", x1, x2)
   x1 = lax.convert_element_type(x1, dtype)
 

diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
@@ -120,6 +120,7 @@ def op_record(name, nargs, dtypes, shapes, rng_factory, diff_modes,
               jtu.rand_default, [], check_dtypes=False),
     op_record("greater", 2, all_dtypes, all_shapes, jtu.rand_some_equal, []),
     op_record("greater_equal", 2, all_dtypes, all_shapes, jtu.rand_some_equal, []),
+    op_record("ldexp", 2, int_dtypes, all_shapes, jtu.rand_default, [], check_dtypes=False),
     op_record("less", 2, all_dtypes, all_shapes, jtu.rand_some_equal, []),
     op_record("less_equal", 2, all_dtypes, all_shapes, jtu.rand_some_equal, []),
     op_record("log", 1, number_dtypes, all_shapes, jtu.rand_positive, ["rev"],