Repository: jax-ml/jax
Pull request: add dtype warnings to array-creation routines
fixes #1230

Unless the `jax_enable_x64` option is set, we truncate all 64 bit dtypes to their 32bit counterparts (and complex128 to complex64) by using the function `xla_bridge.canonicalize_dtype` in both lax_numpy.py and lax.py. According to #1230 that behavior can be surprising when a user writes code that explicitly requests a dtype and it's silently discarded, so we want to produce a warning to help users understand and control what's going on.

I did a quick survey of jax.numpy:
1. `np.array` (of which `np.asarray` is a thin wrapper) calls `canonicalize_dtype` but doesn't use the output in a way that would cause truncation, instead ultimately calling `lax.convert_element_type`
2. `np.ones`, `np.ones_like`, `np.zeros`, `np.zeros_like`, `np.full`, `np.identity`, `np.arange`, `np.linspace`, `np.tri` don't call `canonicalize_dtype` directly and instead call lax functions like `lax.convert_element_type` or `lax.full`

So we could just insert these checks at the lax.py layer rather than the jax.numpy layer, thus requiring less code. But I think it's better to add these checks at the jax.numpy layer because then we can provide the user with information about which jax.numpy function being called is leading to the warning.

This PR makes that change, and also adds pretty thorough checks that we're raising warnings when intended and not raising them otherwise!

Problem statement:
JAX should throw an error (or warning?) when asked to cast to a dtype it does not support on a device
I believe JAX should in some way notify the user when a cast operation cannot be performed.
``` python
a = np.ones((1,2,3))
print(a.dtype) # reports float32
a = np.array(a, dtype='float64') # JAX should throw an error on this line
print(a.dtype) # reports float32
```

Hints from issue discussion:
Great point, you're right! Let's do this.

By the way, if you want to switch on support for 64-bit values (in which case this cast will work as you expect), see the first item in our [Gotchas list](https://github.com/google/jax#current-gotchas).

Commit message: add dtype warnings to array-creation routines

fixes #1230

Patch to review:
diff --git a/jax/lax/__init__.py b/jax/lax/__init__.py
index 26d6a8274a83..47eabb8f81d3 100644
--- a/jax/lax/__init__.py
+++ b/jax/lax/__init__.py
@@ -18,7 +18,7 @@
                   _reduce_and, _reduce_window_sum, _reduce_window_max,
                   _reduce_window_min, _reduce_window_prod, _float, _complex,
                   _input_dtype, _const, _eq_meet, _safe_mul,
-                  _broadcasting_select)
+                  _broadcasting_select, _check_user_dtype_supported)
 from .lax_control_flow import *
 from .lax_fft import *
 from .lax_parallel import *
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
index 5f9b58465fb9..dc314ce4fae2 100644
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -4408,3 +4408,15 @@ def subvals(lst, replace):
 
 def _abstractify(x):
   return raise_to_shaped(core.get_aval(x))
+
+
+def _check_user_dtype_supported(dtype, fun_name=None):
+  if dtype is not None and onp.dtype(dtype) != xla_bridge.canonicalize_dtype(dtype):
+    msg = ("Explicitly requested dtype {} {} is not available, "
+           "and will be truncated to dtype {}. To enable more dtypes, set the "
+           "jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell "
+           "environment variable. "
+           "See https://github.com/google/jax#current-gotchas for more.")
+    fun_name = "requested in {}".format(fun_name) if fun_name else ""
+    truncated_dtype = xla_bridge.canonicalize_dtype(dtype).name
+    warnings.warn(msg.format(dtype, fun_name , truncated_dtype))
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
index 6f871f1c9c94..dfacf22072ab 100644
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -1408,6 +1408,7 @@ def atleast_3d(*arys):
 def array(object, dtype=None, copy=True, order="K", ndmin=0):
   if order is not None and order != "K":
     raise NotImplementedError("Only implemented for order='K'")
+  lax._check_user_dtype_supported(dtype, "array")
 
   if isinstance(object, ndarray):
     if dtype and _dtype(object) != xla_bridge.canonicalize_dtype(dtype):
@@ -1442,38 +1443,49 @@ def array(object, dtype=None, copy=True, order="K", ndmin=0):
 
 @_wraps(onp.asarray)
 def asarray(a, dtype=None, order=None):
+  lax._check_user_dtype_supported(dtype, "asarray")
   return array(a, dtype=dtype, copy=False, order=order)
 
 
 @_wraps(onp.zeros_like)
 def zeros_like(x, dtype=None):
+  lax._check_user_dtype_supported(dtype, "zeros_like")
   return lax.full_like(x, 0, dtype)
 
 
 @_wraps(onp.ones_like)
 def ones_like(x, dtype=None):
+  lax._check_user_dtype_supported(dtype, "ones_like")
   return lax.full_like(x, 1, dtype)
 
 
 @_wraps(onp.full)
 def full(shape, fill_value, dtype=None):
+  lax._check_user_dtype_supported(dtype, "full")
   return lax.full(shape, fill_value, dtype)
 
 
 @_wraps(onp.full_like)
 def full_like(a, fill_value, dtype=None):
+  lax._check_user_dtype_supported(dtype, "full_like")
   return lax.full_like(a, fill_value, dtype)
 
 
 @_wraps(onp.zeros)
-def zeros(shape, dtype=onp.dtype("float64")):
+def zeros(shape, dtype=None):
   if isinstance(shape, types.GeneratorType):
     raise TypeError("expected sequence object with len >= 0 or a single integer")
+  dtype = onp.dtype("float64") if dtype is None else dtype
+  lax._check_user_dtype_supported(dtype, "zeros")
   shape = (shape,) if onp.isscalar(shape) else shape
   return lax.full(shape, 0, dtype)
 
 @_wraps(onp.ones)
-def ones(shape, dtype=onp.dtype("float64")):
+def ones(shape, dtype=None):
+  if isinstance(shape, types.GeneratorType):
+    raise TypeError("expected sequence object with len >= 0 or a single integer")
+  lax._check_user_dtype_supported(dtype, "ones")
+  dtype = onp.dtype("float64") if dtype is None else dtype
   shape = (shape,) if onp.isscalar(shape) else shape
   return lax.full(shape, 1, dtype)
 
@@ -1493,7 +1505,9 @@ def array_equal(a1, a2):
 
 
 @_wraps(onp.eye)
-def eye(N, M=None, k=None, dtype=onp.dtype("float64")):
+def eye(N, M=None, k=None, dtype=None):
+  lax._check_user_dtype_supported(dtype, "eye")
+  dtype = onp.dtype("float64") if dtype is None else dtype
   M = N if M is None else M
   if N < 0 or M < 0:
     msg = "negative dimensions are not allowed, got {} and {}"
@@ -1512,11 +1526,13 @@ def eye(N, M=None, k=None, dtype=onp.dtype("float64")):
 
 @_wraps(onp.identity)
 def identity(n, dtype=None):
+  lax._check_user_dtype_supported(dtype, "identity")
   return eye(n, dtype=dtype)
 
 
 @_wraps(onp.arange)
 def arange(start, stop=None, step=None, dtype=None):
+  lax._check_user_dtype_supported(dtype, "arange")
   # If called like np.arange(N), we create a lazy lax._IotaConstant.
   if stop is None and step is None:
     dtype = dtype or _dtype(start)
@@ -1538,6 +1554,7 @@ def wrapper(*args, **kwargs):
 
 def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None,
              axis=0):
+  lax._check_user_dtype_supported(dtype, "linspace")
   try:
     out = onp.linspace(start, stop, num, endpoint, retstep, dtype, axis)
     if retstep:
@@ -1646,6 +1663,7 @@ def repeat(a, repeats, axis=None):
 
 @_wraps(onp.tri)
 def tri(N, M=None, k=0, dtype=None):
+  lax._check_user_dtype_supported(dtype, "tri")
   M = M if M is not None else N
   dtype = dtype or float32
   x = arange(N, dtype=int32)
@@ -1679,6 +1697,7 @@ def triu(m, k=0):
 def trace(a, offset=0, axis1=0, axis2=1, dtype=None, out=None):
   if out:
     raise NotImplementedError("The 'out' argument to trace is not supported.")
+  lax._check_user_dtype_supported(dtype, "trace")
 
   axis1 = _canonicalize_axis(axis1, ndim(a))
   axis2 = _canonicalize_axis(axis2, ndim(a))
@@ -2839,6 +2858,10 @@ def median(a, axis=None, out=None, overwrite_input=False, keepdims=False):
     return quantile(a, q, axis=axis, out=out, overwrite_input=overwrite_input,
                     keepdims=keepdims)
 
+def _astype(arr, dtype):
+  lax._check_user_dtype_supported(dtype, "astype")
+  return lax.convert_element_type(arr, dtype)
+
 ### track unimplemented functions
 
 def _not_implemented(fun):
@@ -2934,7 +2957,7 @@ def _unimplemented_setitem(self, i, x):
 setattr(ShapedArray, "T", core.aval_property(transpose))
 setattr(ShapedArray, "real", core.aval_property(real))
 setattr(ShapedArray, "imag", core.aval_property(imag))
-setattr(ShapedArray, "astype", core.aval_method(lax.convert_element_type))
+setattr(ShapedArray, "astype", core.aval_method(_astype))
 
 
 # Forward operators, methods, and properties on DeviceArray to lax_numpy
@@ -2948,7 +2971,7 @@ def _unimplemented_setitem(self, i, x):
 setattr(DeviceArray, "T", property(transpose))
 setattr(DeviceArray, "real", property(real))
 setattr(DeviceArray, "imag", property(imag))
-setattr(DeviceArray, "astype", lax.convert_element_type)
+setattr(DeviceArray, "astype", _astype)
 
 
 # Extra methods that are handy
diff --git a/tests/api_test.py b/tests/api_test.py
index a49e205df16d..8a5affa9843a 100644
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -19,6 +19,7 @@
 import collections
 from functools import partial
 import unittest
+import warnings
 
 from absl.testing import absltest
 import numpy as onp
@@ -41,6 +42,7 @@
 
 from jax.config import config
 config.parse_flags_with_absl()
+FLAGS = config.FLAGS
 
 class APITest(jtu.JaxTestCase):
 
@@ -975,6 +977,53 @@ def f(x):
     for x, y in zip(xs, ys):
       self.assertAllClose(x * 2 - 3., y, check_dtypes=True)
 
+  def test_issue_1230(self):
+    if FLAGS.jax_enable_x64:
+      return  # test only applies when x64 is disabled
+
+    def check_warning(warn, nowarn):
+      with warnings.catch_warnings(record=True) as w:
+        warnings.simplefilter("always")
+
+        nowarn()  # get rid of extra startup warning
+
+        prev_len = len(w)
+        nowarn()
+        assert len(w) == prev_len
+
+        warn()
+        assert len(w) > 0
+        msg = str(w[-1].message)
+        expected_prefix = "Explicitly requested dtype "
+        self.assertEqual(expected_prefix, msg[:len(expected_prefix)])
+
+        prev_len = len(w)
+        nowarn()
+        assert len(w) == prev_len
+
+    check_warning(lambda: np.array([1, 2, 3], dtype="float64"),
+                  lambda: np.array([1, 2, 3], dtype="float32"),)
+    check_warning(lambda: np.ones(3, dtype=onp.float64),
+                  lambda: np.ones(3))
+    check_warning(lambda: np.ones_like(3, dtype=onp.int64),
+                  lambda: np.ones_like(3, dtype=onp.int32))
+    check_warning(lambda: np.zeros(3, dtype="int64"),
+                  lambda: np.zeros(3, dtype="int32"))
+    check_warning(lambda: np.zeros_like(3, dtype="float64"),
+                  lambda: np.zeros_like(3, dtype="float32"))
+    check_warning(lambda: np.full((2, 3), 1, dtype="int64"),
+                  lambda: np.full((2, 3), 1))
+    check_warning(lambda: np.ones(3).astype("float64"),
+                  lambda: np.ones(3).astype("float32"))
+    check_warning(lambda: np.eye(3, dtype=onp.float64),
+                  lambda: np.eye(3))
+    check_warning(lambda: np.arange(3, dtype=onp.float64),
+                  lambda: np.arange(3, dtype=onp.float32))
+    check_warning(lambda: np.linspace(0, 3, dtype=onp.float64),
+                  lambda: np.linspace(0, 3, dtype=onp.float32))
+    check_warning(lambda: np.tri(2, dtype="float64"),
+                  lambda: np.tri(2, dtype="float32"))
+
 
 if __name__ == '__main__':
   absltest.main()


Code review comments:
diff_hunk: @@ -1442,38 +1443,49 @@ def array(object, dtype=None, copy=True, order="K", ndmin=0):
 
 @_wraps(onp.asarray)
 def asarray(a, dtype=None, order=None):
+  lax._check_user_dtype_supported(dtype, "asarray")
   return array(a, dtype=dtype, copy=False, order=order)
 
 
 @_wraps(onp.zeros_like)
 def zeros_like(x, dtype=None):
+  lax._check_user_dtype_supported(dtype, "zeros_like")
   return lax.full_like(x, 0, dtype)
 
 
 @_wraps(onp.ones_like)
 def ones_like(x, dtype=None):
+  lax._check_user_dtype_supported(dtype, "ones_like")
   return lax.full_like(x, 1, dtype)
 
 
 @_wraps(onp.full)
 def full(shape, fill_value, dtype=None):
+  lax._check_user_dtype_supported(dtype, "full")
   return lax.full(shape, fill_value, dtype)
 
 
 @_wraps(onp.full_like)
 def full_like(a, fill_value, dtype=None):
+  lax._check_user_dtype_supported(dtype, "full_like")
   return lax.full_like(a, fill_value, dtype)
 
 
 @_wraps(onp.zeros)
-def zeros(shape, dtype=onp.dtype("float64")):
+def zeros(shape, dtype=None):
   if isinstance(shape, types.GeneratorType):
     raise TypeError("expected sequence object with len >= 0 or a single integer")
+  dtype = onp.dtype("float64") if dtype is None else dtype
+  lax._check_user_dtype_supported(dtype, "zeros")
path: jax/numpy/lax_numpy.py
text: @user1:
This should go above setting dtype right?

@author:
Excellent catch, thanks!