Repository: google/jax

Issue:
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
```


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).

Resolving patch:
diff --git a/jax/lax/__init__.py b/jax/lax/__init__.py
--- 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
--- 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
--- 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")
+  lax._check_user_dtype_supported(dtype, "zeros")
+  dtype = onp.dtype("float64") if dtype is None else dtype
   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


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- 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,54 @@ def f(x):
     for x, y in zip(xs, ys):
       self.assertAllClose(x * 2 - 3., y, check_dtypes=True)
 
+  def test_dtype_warning(self):
+    # cf. issue #1230
+    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()