Repository: dask/dask

Issue:
da.core.elemwise undefined dtype inference on 0d Arrays
I've noticed an undefined behaviour problem with the dtype inference in elemwise. Here's the relevant bit of code:
```python
        vals = [np.empty((1,) * a.ndim, dtype=a.dtype)
                if not is_scalar_for_elemwise(a) else a
                for a in args]
        dt = apply_infer_dtype(op, vals, {}, 'elemwise', suggest_dtype=False)
```
If one of the arguments is a 0d dask array, then apply_infer_dtype is instead given an empty 0d numpy array (as the comment immediately above this code indicates, 0d dask arrays are not considered scalar for elemwise).

The problem is that a 0d numpy array is treated the same as a scalar by numpy, and hence the *value* affects dtype inference. Because `np.empty` was used, the value is undefined.

The simple solution is to replace `np.empty` with `np.ones`, but I wanted to raise the question of what the dtype *should* be in this case. By doing so, the dtypes of scalars is effectively ignored (because the value 1 won't cause any promotions), which makes it "optimistic" (assumes the true value won't cause overflows), whereas perhaps it ought to be conservative (large scalar types to cause promotions). For example, just changing `np.empty` to `np.ones` will give
```python
>>> a = da.ones(1, dtype=np.float32, chunks=1)
>>> b = da.from_array(np.array(1e200), chunks=())
>>> (a + b).dtype
dtype('float32')
```
even though the result doesn't fit into a finite float32 because b is too large (there is a separate issue that `(a + b).compute()` will actually return a float64 numpy array, but I'll file a separate issue for that when I get around to it).


Discussion:

Resolving patch:
diff --git a/dask/array/__init__.py b/dask/array/__init__.py
--- a/dask/array/__init__.py
+++ b/dask/array/__init__.py
@@ -9,7 +9,8 @@
                        bincount, digitize, histogram, cov, array, dstack,
                        vstack, hstack, compress, extract, round, count_nonzero,
                        flatnonzero, nonzero, around, isnull, notnull, isclose,
-                       corrcoef, swapaxes, tensordot, transpose, dot)
+                       corrcoef, swapaxes, tensordot, transpose, dot,
+                       result_type)
 from .reshape import reshape
 from .ufunc import (add, subtract, multiply, divide, logaddexp, logaddexp2,
         true_divide, floor_divide, negative, power, remainder, mod, conj, exp,
diff --git a/dask/array/core.py b/dask/array/core.py
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -2488,6 +2488,7 @@ def elemwise(op, *args, **kwargs):
     arrays = [a for a in args if not is_scalar_for_elemwise(a)]
     other = [(i, a) for i, a in enumerate(args) if is_scalar_for_elemwise(a)]
 
+    need_enforce_dtype = False
     if 'dtype' in kwargs:
         dt = kwargs['dtype']
     else:
@@ -2496,24 +2497,30 @@ def elemwise(op, *args, **kwargs):
         # their values to compute the result dtype:
         # https://github.com/numpy/numpy/issues/6240
         # We don't inspect the values of 0d dask arrays, because these could
-        # hold potentially very expensive calculations.
-        vals = [np.empty((1,) * a.ndim, dtype=a.dtype)
+        # hold potentially very expensive calculations. Instead, we treat
+        # them just like other arrays, and if necessary cast the result of op
+        # to match.
+        vals = [np.empty((1,) * max(1, a.ndim), dtype=a.dtype)
                 if not is_scalar_for_elemwise(a) else a
                 for a in args]
         dt = apply_infer_dtype(op, vals, {}, 'elemwise', suggest_dtype=False)
+        need_enforce_dtype = any(not is_scalar_for_elemwise(a) and a.ndim == 0 for a in args)
 
     name = kwargs.get('name', None) or '%s-%s' % (funcname(op),
                                                   tokenize(op, dt, *args))
 
+    atop_kwargs = dict(dtype=dt, name=name, token=funcname(op).strip('_'))
     if other:
-        result = atop(partial_by_order, expr_inds,
-                      *concat((a, tuple(range(a.ndim)[::-1])) for a in arrays),
-                      dtype=dt, name=name, function=op, other=other,
-                      token=funcname(op))
-    else:
-        result = atop(op, expr_inds,
-                      *concat((a, tuple(range(a.ndim)[::-1])) for a in arrays),
-                      dtype=dt, name=name)
+        atop_kwargs['function'] = op
+        atop_kwargs['other'] = other
+        op = partial_by_order
+    if need_enforce_dtype:
+        atop_kwargs['enforce_dtype'] = dt
+        atop_kwargs['enforce_dtype_function'] = op
+        op = _enforce_dtype
+    result = atop(op, expr_inds,
+                  *concat((a, tuple(range(a.ndim)[::-1])) for a in arrays),
+                  **atop_kwargs)
 
     return handle_out(out, result)
 
@@ -2548,6 +2555,37 @@ def handle_out(out, result):
         return result
 
 
+def _enforce_dtype(*args, **kwargs):
+    """Calls a function and converts its result to the given dtype.
+
+    The parameters have deliberately been given unwieldy names to avoid
+    clashes with keyword arguments consumed by atop or partial_by_order.
+
+    A dtype of `object` is treated as a special case and not enforced,
+    because it is used as a dummy value in some places when the result will
+    not be a block in an Array.
+
+    Parameters
+    ----------
+    enforce_dtype : dtype
+        Result dtype
+    enforce_dtype_function : callable
+        The wrapped function, which will be passed the remaining arguments
+    """
+    dtype = kwargs.pop('enforce_dtype')
+    function = kwargs.pop('enforce_dtype_function')
+
+    result = function(*args, **kwargs)
+    if dtype != object:
+        if np.isscalar(result):
+            # scalar astype method doesn't take the keyword arguments, so
+            # have to convert via 0-dimensional array and back.
+            result = result[...].astype(dtype, casting='same_kind', copy=False)[()]
+        else:
+            result = result.astype(dtype, casting='same_kind', copy=False)
+    return result
+
+
 @wraps(chunk.broadcast_to)
 def broadcast_to(x, shape):
     x = asarray(x)


Test patch:
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -395,15 +395,19 @@ def test_broadcast_shapes():
 
 
 def test_elemwise_on_scalars():
-    x = np.arange(10)
+    x = np.arange(10, dtype=np.int64)
     a = from_array(x, chunks=(5,))
     assert len(a._keys()) == 2
     assert_eq(a.sum()**2, x.sum()**2)
 
-    x = np.arange(11)
-    a = from_array(x, chunks=(5,))
-    assert len(a._keys()) == 3
-    assert_eq(a, x)
+    y = np.arange(10, dtype=np.int32)
+    b = from_array(y, chunks=(5,))
+    result = a.sum() * b
+    # Dask 0-d arrays do not behave like numpy scalars for type promotion
+    assert result.dtype == np.int64
+    assert result.compute().dtype == np.int64
+    assert (x.sum() * y).dtype == np.int32
+    assert_eq((x.sum() * y).astype(np.int64), result)
 
 
 def test_partial_by_order():
@@ -438,6 +442,20 @@ def test_elemwise_differently_chunked():
     assert_eq(b + a, x + y)
 
 
+def test_elemwise_dtype():
+    values = [
+        da.from_array(np.ones(5, np.float32), chunks=3),
+        da.from_array(np.ones(5, np.int16), chunks=3),
+        da.from_array(np.ones(5, np.int64), chunks=3),
+        da.from_array(np.ones((), np.float64), chunks=()) * 1e200,
+        np.ones(5, np.float32),
+        1, 1.0, 1e200, np.int64(1), np.ones((), np.int64),
+    ]
+    for x in values:
+        for y in values:
+            assert da.maximum(x, y).dtype == da.result_type(x, y)
+
+
 def test_operators():
     x = np.arange(10)
     y = np.arange(10).reshape((10, 1))
diff --git a/dask/array/tests/test_routines.py b/dask/array/tests/test_routines.py
--- a/dask/array/tests/test_routines.py
+++ b/dask/array/tests/test_routines.py
@@ -799,3 +799,24 @@ def test_multi_insert():
     c = da.from_array(z, chunks=(1, 2))
     assert_eq(np.insert(np.insert(z, [0, 1], -1, axis=0), [1], -1, axis=1),
               da.insert(da.insert(c, [0, 1], -1, axis=0), [1], -1, axis=1))
+
+
+def test_result_type():
+    a = da.from_array(np.ones(5, np.float32), chunks=(3,))
+    b = da.from_array(np.ones(5, np.int16), chunks=(3,))
+    c = da.from_array(np.ones(5, np.int64), chunks=(3,))
+    x = np.ones(5, np.float32)
+    assert da.result_type(b, c) == np.int64
+    assert da.result_type(a, b, c) == np.float64
+    assert da.result_type(b, np.float32) == np.float32
+    assert da.result_type(b, np.dtype(np.float32)) == np.float32
+    assert da.result_type(b, x) == np.float32
+    # Effect of scalars depends on their value
+    assert da.result_type(1, b) == np.int16
+    assert da.result_type(1.0, a) == np.float32
+    assert da.result_type(np.int64(1), b) == np.int16
+    assert da.result_type(np.ones((), np.int64), b) == np.int16  # 0d array
+    assert da.result_type(1e200, a) == np.float64   # 1e200 is too big for float32
+    # dask 0d-arrays are NOT treated like scalars
+    c = da.from_array(np.ones((), np.float64), chunks=())
+    assert da.result_type(a, c) == np.float64