Repository: google/jax

Issue:
Potential type promotion bug.
Since the latest jax type promotion changes I'm seeing the following strange type promotion behavior (x64, CPU at head):

```python
import jax.numpy as np
from jax.config import config
config.update('jax_enable_x64', True)
xs = np.array([1., 2.], np.float32)
sigma = np.float32(1.0)
print(np.exp(xs / sigma ** 2).dtype)  # Float64
```


Discussion:
This isn't really a bug, but I agree it's something we might like to improve. I believe this behavior hasn't been altered by the recent type promotion changes; it has always held.

This behavior follows from two facts:
* `jax.numpy.float32` is an alias for `numpy.float32`. This means it acts identically in every way.
* Numpy has the following promotion behavior:
```
In [6]: type(onp.float32(1.0) ** 2)
Out[6]: numpy.float64
```
i.e., the RHS argument to the `/` operator in your example above is a `float64`, which JAX respects. So this is really a NumPy vs JAX difference.

The best fixes I can think of are:
a) use an array instead of a `NumPy` scalar:
```
print(np.exp(xs / np.array(sigma) ** 2).dtype)
float32
```
b) we change JAX to offer its own scalar objects whose promotion behavior matches the JAX array promotion behavior.
c) we change `jax.numpy.float32(x)` to return a JAX array object.
For what it's worth, I think the idiomatic way to write something like this, in both JAX and NumPy, is now to use Python floats:
```python
xs = np.array([1., 2.], np.float32)
sigma = 1.0
print(np.exp(xs / sigma ** 2).dtype)  # float32
```
These end up giving the same result in `jax.numpy` and `numpy`, but for two different reasons:
- JAX special cases Python's builtin `float` in its casting rules as having lower priority than array dtypes.
- NumPy has *value dependent casting* for scalars and 0d arrays. So NumPy has different behavior than JAX if you use Python scalars not representable in float32, e.g., `sigma = 1e40`.

> c) we change `jax.numpy.float32(x)` to return a JAX array object.

This seems like a reasonable choice, but keep in mind we will probably also have to special-case `jax.numpy.float32` (whatever it is) when it's passed into a `dtype` argument in JAX. One of the most common uses for these type scalars is to write something like `x.astype(np.float32)`.

The downside is that there is no way we can make `x.astype(np.float32)`  work if `x` is a NumPy array. Users will get `TypeError: data type not understood`.

Given that it may be impossible to make something that matches all user expectations, this may be an argument for dropping aliases like `jax.numpy.float32` entirely, though I'm sure that would encounter resistance :)
@shoyer I'm not so sure about that:
```
In [1]: import numpy as onp; from jax import numpy as np, lax; import jax                             
c
In [2]: class Foo(onp.int32): 
   ...:     pass 
   ...:                  

In [3]: Foo              
Out[3]: __main__.Foo

In [4]: Foo(3)           
Out[4]: 3

In [5]: onp.array([2, 4], dtype=Foo)               
Out[5]: array([2, 4], dtype=int32)

In [8]: x = onp.random.randn(3, 4)                                                                        

In [9]: x                                                                                                 
Out[9]: 
array([[ 0.61102705, -1.28747357, -0.87507959, -0.83302679],
       [-1.21343036, -1.44278927,  1.71148559,  1.00484012],
       [-0.09954842, -0.32945393,  0.75188577,  1.38867382]])

In [10]: x.astype(Foo)                                                                                    
Out[10]: 
array([[ 0, -1,  0,  0],
       [-1, -1,  1,  1],
       [ 0,  0,  0,  1]], dtype=int32)

```
I don't know why it works, but it does!
```
import numpy as onp
from jax import numpy as np

class MyMeta(type):
  def __call__(cls, *args, **kwds):
    return np.array(*args, dtype=onp.int32)

class MyInt32(onp.int32, metaclass=MyMeta):
  pass
```
seems pretty promising.

```
In [31]: MyInt32(77)                                                                                      
Out[31]: DeviceArray(77, dtype=int32)
```
This looks like good luck rather than an intentional feature, so I would be hesitant to count on long term :)

Either way, if we use it definitely use `__new__` rather than writing a meta class.
This appears not to be an accident:
https://github.com/numpy/numpy/blob/bf9614b3253e9ef4dbbcef1f48751955b8f2a598/numpy/core/src/multiarray/scalarapi.c#L466

To quote:
```
    /*
     * Otherwise --- type is a sub-type of an array scalar
     * not corresponding to a registered data-type object.
     */
```

The numpy tests also pointed me to what appears to be a slightly cleaner solution:
https://github.com/numpy/numpy/blob/bf9614b3253e9ef4dbbcef1f48751955b8f2a598/numpy/core/tests/test_dtype.py#L1057

```
import numpy as onp
from jax import numpy as np

class MyInt32(object):
  dtype = onp.dtype(onp.int32)
  def __new__(cls, x):
    return np.array(x, dtype=onp.int32)
```

i.e., it appears to be a tested property of `np.dtype` that it checks for a `dtype` attribute.
> i.e., it appears to be a tested property of `np.dtype` that it checks for a `dtype` attribute.

Yes, this looks like a clean solution. It's tested, so we can probably rely on it, though of course it appears to be undocumented!

Resolving patch:
diff --git a/jax/dtypes.py b/jax/dtypes.py
--- a/jax/dtypes.py
+++ b/jax/dtypes.py
@@ -135,7 +135,14 @@ def finfo(dtype):
 
 def issubdtype(a, b):
   if a == bfloat16:
-    return b in [onp.floating, onp.inexact, onp.number]
+    return b in [bfloat16, _bfloat16_dtype, onp.floating, onp.inexact,
+                 onp.number]
+  if not issubclass(b, onp.generic):
+    # Workaround for JAX scalar types. NumPy's issubdtype has a backward
+    # compatibility behavior for the second argument of issubdtype that
+    # interacts badly with JAX's custom scalar types. As a workaround,
+    # explicitly cast the second argument to a NumPy type object.
+    b = onp.dtype(b).type
   return onp.issubdtype(a, b)
 
 can_cast = onp.can_cast
diff --git a/jax/lax_linalg.py b/jax/lax_linalg.py
--- a/jax/lax_linalg.py
+++ b/jax/lax_linalg.py
@@ -94,7 +94,8 @@ def g(c, *args, **kwargs):
 
 # primitives
 
-_cpu_lapack_types = {np.float32, np.float64, np.complex64, np.complex128}
+_cpu_lapack_types = {onp.dtype(onp.float32), onp.dtype(onp.float64),
+                     onp.dtype(onp.complex64), onp.dtype(onp.complex128)}
 
 # Cholesky decomposition
 
@@ -137,7 +138,7 @@ def _nan_like(c, operand):
 def cholesky_cpu_translation_rule(c, operand):
   shape = c.GetShape(operand)
   dtype = shape.element_type().type
-  if len(shape.dimensions()) == 2 and dtype in _cpu_lapack_types:
+  if len(shape.dimensions()) == 2 and onp.dtype(dtype) in _cpu_lapack_types:
     result, info = lapack.potrf(c, operand, lower=True)
     return c.Select(c.Eq(info, c.ConstantS32Scalar(0)), result,
                     _nan_like(c, result))
@@ -392,7 +393,7 @@ def _triangular_solve_cpu_translation_rule(
     c, a, b, left_side, lower, transpose_a, conjugate_a, unit_diagonal):
   shape = c.GetShape(a)
   dtype = shape.element_type().type
-  if len(shape.dimensions()) == 2 and dtype in _cpu_lapack_types:
+  if len(shape.dimensions()) == 2 and onp.dtype(dtype) in _cpu_lapack_types:
     if conjugate_a and not transpose_a:
       a = c.Conj(a)
       conjugate_a = False
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
@@ -126,29 +126,50 @@ def __init__(shape, dtype=None, buffer=None, offset=0, strides=None,
 size = onp.size
 _dtype = dtypes.result_type
 
-bool_ = onp.bool_
-int_ = dtypes.int_
-float_ = dtypes.float_
-complex_ = dtypes.complex_
-
-uint8 = onp.uint8
-uint16 = onp.uint16
-uint32 = onp.uint32
-uint64 = onp.uint64
-int8 = onp.int8
-int16 = onp.int16
-int32 = onp.int32
-int64 = onp.int64
-bfloat16 = dtypes.bfloat16
-float16 = onp.float16
-float32 = single = onp.float32
-float64 = double = onp.float64
-complex64 = csingle = onp.complex64
-complex128 = cdouble = onp.complex128
+# At present JAX doesn't have a reason to distinguish between scalars and arrays
+# in its object system. Further, we want JAX scalars to have the same type
+# promotion behaviors as JAX arrays. Rather than introducing a new type of JAX
+# scalar object with JAX promotion behaviors, instead we make the JAX scalar
+# types return JAX arrays when instantiated.
+
+class _ScalarMeta(type):
+  def __hash__(self):
+    return hash(self.dtype.type)
+
+  def __eq__(self, other):
+    return id(self) == id(other) or self.dtype == other
+
+  def __ne__(self, other):
+    return not (self == other)
+
+  def __call__(self, x):
+    return array(self.dtype.type(x), dtype=self.dtype)
+
+def _make_scalar_type(onp_scalar_type):
+  return type(onp_scalar_type.__name__,
+              (six.with_metaclass(_ScalarMeta, object),),
+              {"dtype": onp.dtype(onp_scalar_type)})
+
+bool_ = _make_scalar_type(onp.bool_)
+uint8 = _make_scalar_type(onp.uint8)
+uint16 = _make_scalar_type(onp.uint16)
+uint32 = _make_scalar_type(onp.uint32)
+uint64 = _make_scalar_type(onp.uint64)
+int8 = _make_scalar_type(onp.int8)
+int16 = _make_scalar_type(onp.int16)
+int32 = _make_scalar_type(onp.int32)
+int64 = _make_scalar_type(onp.int64)
+bfloat16 = _make_scalar_type(dtypes.bfloat16)
+float16 = _make_scalar_type(onp.float16)
+float32 = single = _make_scalar_type(onp.float32)
+float64 = double = _make_scalar_type(onp.float64)
+complex64 = csingle = _make_scalar_type(onp.complex64)
+complex128 = cdouble = _make_scalar_type(onp.complex128)
+
+int_ = int32 if dtypes.int_ == onp.int32 else int64
+float_ = float32 if dtypes.float_ == onp.float32 else float64
+complex_ = complex64 if dtypes.complex_ == onp.complex64 else complex128
 
-flexible = onp.flexible
-character = onp.character
-object_ = onp.object_
 number = onp.number
 inexact = onp.inexact
 complexfloating = onp.complexfloating
@@ -157,6 +178,10 @@ def __init__(shape, dtype=None, buffer=None, offset=0, strides=None,
 signedinteger = onp.signedinteger
 unsignedinteger = onp.unsignedinteger
 
+flexible = onp.flexible
+character = onp.character
+object_ = onp.object_
+
 iinfo = dtypes.iinfo
 
 dtype = onp.dtype
@@ -363,8 +388,6 @@ def finfo(dtype): return dtypes.finfo(dtype)
 @_wraps(onp.issubdtype)
 def issubdtype(arg1, arg2): return dtypes.issubdtype(arg1, arg2)
 
-issubdtype = dtypes.issubdtype
-
 @_wraps(onp.isscalar)
 def isscalar(num): return dtypes.is_python_scalar(num) or onp.isscalar(num)
 
@@ -459,7 +482,7 @@ def _logical_op(np_op, bitwise_op):
   @_wraps(np_op, update_doc=False)
   def op(*args):
     zero = lambda x: lax.full_like(x, shape=(), fill_value=0)
-    args = (x if issubdtype(_dtype(x), onp.bool_) else lax.ne(x, zero(x))
+    args = (x if issubdtype(_dtype(x), bool_) else lax.ne(x, zero(x))
             for x in args)
     return bitwise_op(*_promote_args(np_op.__name__, *args))
   return op
@@ -481,7 +504,7 @@ def divide(x1, x2):
   # decide whether to perform integer division based on Numpy result dtype, as a
   # way to check whether Python 3 style division is active in Numpy
   result_dtype = _result_dtype(onp.divide, x1, x2)
-  if issubdtype(result_dtype, onp.integer):
+  if issubdtype(result_dtype, integer):
     return floor_divide(x1, x2)
   else:
     return true_divide(x1, x2)
@@ -514,7 +537,7 @@ def floor_divide(x1, x2):
 @_wraps(onp.divmod)
 def divmod(x1, x2):
   x1, x2 = _promote_args("divmod", x1, x2)
-  if issubdtype(_dtype(x1), onp.integer):
+  if issubdtype(_dtype(x1), integer):
     return floor_divide(x1, x2), remainder(x1, x2)
   else:
     return _float_divmod(x1, x2)
@@ -704,7 +727,7 @@ def arcsinh(x):
   x, = _promote_dtypes_inexact(x)
   one = lax._const(x, 1)
   result = lax.log(x + lax.sqrt(x * x + one))
-  if issubdtype(_dtype(result), onp.complexfloating):
+  if issubdtype(_dtype(result), complexfloating):
     return result
   a = abs(x)
   sqrt_max_value = onp.sqrt(finfo(_dtype(x)).max)
@@ -723,7 +746,7 @@ def arccosh(x):
   x, = _promote_dtypes_inexact(x)
   one = lax._const(x, 1)
   result = lax.log(x + lax.sqrt((x + one) * (x - one)))
-  if issubdtype(_dtype(result), onp.complexfloating):
+  if issubdtype(_dtype(result), complexfloating):
     return result
   sqrt_max_value = onp.sqrt(finfo(_dtype(x)).max)
   log2 = lax._const(x, onp.log(2))
@@ -736,7 +759,7 @@ def arctanh(x):
   x, = _promote_dtypes_inexact(x)
   one = lax._const(x, 1)
   result = lax._const(x, 0.5) * lax.log((one + x) / (one - x))
-  if issubdtype(_dtype(result), onp.complexfloating):
+  if issubdtype(_dtype(result), complexfloating):
     return result
   return lax.select(abs(x) <= 1, result, lax.full_like(x, onp.nan))
 
@@ -981,7 +1004,7 @@ def _maybe_numpy_1_13_isclose_behavior(a, out):
 def where(condition, x=None, y=None):
   if x is None or y is None:
     raise ValueError("Must use the three-argument form of where().")
-  if not issubdtype(_dtype(condition), onp.bool_):
+  if not issubdtype(_dtype(condition), bool_):
     condition = lax.ne(condition, zeros_like(condition))
   x, y = _promote_dtypes(x, y)
   condition, x, y = broadcast_arrays(condition, x, y)
@@ -1070,7 +1093,7 @@ def clip(a, a_min=None, a_max=None):
 
 def _dtype_info(dtype):
   """Helper function for to get dtype info needed for clipping."""
-  if issubdtype(dtype, onp.integer):
+  if issubdtype(dtype, integer):
     return iinfo(dtype)
   return finfo(dtype)
 
@@ -1199,7 +1222,7 @@ def reduction(a, axis=None, dtype=None, out=None, keepdims=False):
       computation_dtype = result_dtype
     a = lax.convert_element_type(a, computation_dtype)
     result = lax.reduce(a, _reduction_init_val(a, init_val),
-                        op if computation_dtype != bool_ else bool_op, dims)
+                        op if computation_dtype != onp.bool_ else bool_op, dims)
     if keepdims:
       shape_with_singletons = lax.subvals(shape(a), zip(dims, (1,) * len(dims)))
       result = lax.reshape(result, shape_with_singletons)
@@ -1224,7 +1247,7 @@ def _reduction_init_val(a, init_val):
   try:
     return onp.array(init_val, dtype=a_dtype)
   except OverflowError:
-    assert issubdtype(a_dtype, onp.integer)
+    assert issubdtype(a_dtype, integer)
     sign, info = onp.sign(init_val), iinfo(a_dtype)
     return onp.array(info.min if sign < 0 else info.max, dtype=a_dtype)
 
@@ -1250,8 +1273,7 @@ def mean(a, axis=None, dtype=None, out=None, keepdims=False):
   else:
     normalizer = onp.prod(onp.take(shape(a), axis))
   if dtype is None:
-    if (issubdtype(_dtype(a), onp.bool_) or
-        issubdtype(_dtype(a), onp.integer)):
+    if issubdtype(_dtype(a), bool_) or issubdtype(_dtype(a), integer):
       dtype = float_
     else:
       dtype = _dtype(a)
@@ -1391,8 +1413,7 @@ def nan_reduction(a, axis=None, out=None, keepdims=False, **kwargs):
 def nanmean(a, axis=None, dtype=None, out=None, keepdims=False):
   if out is not None:
     raise ValueError("nanmean does not support the `out` argument.")
-  if (issubdtype(_dtype(a), onp.bool_) or
-      issubdtype(_dtype(a), onp.integer)):
+  if issubdtype(_dtype(a), bool_) or issubdtype(_dtype(a), integer):
     return mean(a, axis, dtype, out, keepdims)
   if dtype is None:
     dtype = _dtype(a)
@@ -1759,7 +1780,7 @@ def eye(N, M=None, k=None, dtype=None):
     return lax.broadcasted_eye(dtype, (N, M), (0, 1))
   else:
     k_dtype = _dtype(k)
-    if not issubdtype(k_dtype, onp.integer):
+    if not issubdtype(k_dtype, integer):
       msg = "eye argument `k` must be of integer dtype, got {}"
       raise TypeError(msg.format(k_dtype))
     rows = k + lax.broadcasted_iota(k_dtype, (N, M), 0)
@@ -1779,7 +1800,7 @@ def arange(start, stop=None, step=None, dtype=None):
   # 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)
-    if issubdtype(dtype, onp.integer):
+    if issubdtype(dtype, integer):
       return lax.iota(dtype, start)  # avoids materializing
 
   # Fall back to instantiating an ndarray in host memory
@@ -2189,7 +2210,7 @@ def matmul(a, b, precision=None):  # pylint: disable=missing-docstring
 
 @_wraps(onp.vdot, lax_description=_PRECISION_DOC)
 def vdot(a, b, precision=None):
-  if issubdtype(_dtype(a), onp.complexfloating):
+  if issubdtype(_dtype(a), complexfloating):
     a = conj(a)
   return dot(a.ravel(), b.ravel(), precision=precision)
 
@@ -2735,7 +2756,7 @@ def _merge_static_and_dynamic_indices(treedef, static_idx, dynamic_idx):
   return treedef.unflatten(idx)
 
 def _int(aval):
-  return not aval.shape and issubdtype(aval.dtype, onp.integer)
+  return not aval.shape and issubdtype(aval.dtype, integer)
 
 def _index_to_gather(x_shape, idx):
   # Remove ellipses and add trailing slice(None)s.
@@ -2945,8 +2966,8 @@ def _expand_bool_indices(idx):
       abstract_i = core.get_aval(i)
     except TypeError:
       abstract_i = None
-    if (isinstance(abstract_i, ShapedArray) and issubdtype(abstract_i.dtype, onp.bool_)
-          or isinstance(i, list) and _all(not _shape(e) and issubdtype(_dtype(e), onp.bool_)
+    if (isinstance(abstract_i, ShapedArray) and issubdtype(abstract_i.dtype, bool_)
+          or isinstance(i, list) and _all(not _shape(e) and issubdtype(_dtype(e), bool_)
                                           for e in i)):
       if isinstance(i, list):
         i = array(i)


Test patch:
diff --git a/tests/dtypes_test.py b/tests/dtypes_test.py
--- a/tests/dtypes_test.py
+++ b/tests/dtypes_test.py
@@ -34,6 +34,7 @@
 from jax import dtypes
 from jax import numpy as np
 from jax import test_util as jtu
+from jax.interpreters import xla
 
 from jax.config import config
 config.parse_flags_with_absl()
@@ -58,6 +59,10 @@
 all_dtypes = (bool_dtypes + signed_dtypes + unsigned_dtypes + float_dtypes +
               complex_dtypes)
 
+scalar_types = [np.bool_, np.int8, np.int16, np.int32, np.int64,
+                np.uint8, np.uint16, np.uint32, np.uint64,
+                np.bfloat16, np.float16, np.float32, np.float64,
+                np.complex64, np.complex128]
 
 class DtypesTest(jtu.JaxTestCase):
 
@@ -138,6 +143,31 @@ def testPromoteDtypes(self):
           self.assertEqual(onp.promote_types(t1, t2),
                            dtypes.promote_types(t1, t2))
 
+  def testScalarInstantiation(self):
+    for t in [np.bool_, np.int32, np.bfloat16, np.float32, np.complex64]:
+      a = t(1)
+      self.assertEqual(a.dtype, np.dtype(t))
+      self.assertIsInstance(a, xla.DeviceArray)
+      self.assertEqual(0, np.ndim(a))
+
+  def testIsSubdtype(self):
+    for t in scalar_types:
+      self.assertTrue(dtypes.issubdtype(t, t))
+      self.assertTrue(dtypes.issubdtype(onp.dtype(t).type, t))
+      self.assertTrue(dtypes.issubdtype(t, onp.dtype(t).type))
+      if t != np.bfloat16:
+        for category in [onp.generic, np.inexact, np.integer, np.signedinteger,
+                         np.unsignedinteger, np.floating, np.complexfloating]:
+          self.assertEqual(dtypes.issubdtype(t, category),
+                           onp.issubdtype(onp.dtype(t).type, category))
+          self.assertEqual(dtypes.issubdtype(t, category),
+                           onp.issubdtype(onp.dtype(t).type, category))
+
+  def testArrayCasts(self):
+    for t in [np.bool_, np.int32, np.bfloat16, np.float32, np.complex64]:
+      a = onp.array([1, 2.5, -3.7])
+      self.assertEqual(a.astype(t).dtype, np.dtype(t))
+      self.assertEqual(np.array(a).astype(t).dtype, np.dtype(t))
 
   @unittest.skipIf(six.PY2, "Test requires Python 3")
   def testEnumPromotion(self):
diff --git a/tests/lax_control_flow_test.py b/tests/lax_control_flow_test.py
--- a/tests/lax_control_flow_test.py
+++ b/tests/lax_control_flow_test.py
@@ -329,7 +329,7 @@ def testForiLoopErrors(self):
     """Test typing error messages for while."""
     with self.assertRaisesRegex(
       TypeError, "arguments to fori_loop must have equal types"):
-      lax.fori_loop(np.int16(0), np.int32(10), (lambda i, c: c), np.float32(7))
+      lax.fori_loop(onp.int16(0), np.int32(10), (lambda i, c: c), np.float32(7))
 
   def testForiLoopBatched(self):
     def body_fun(i, loop_carry):