Repository: numba/numba

Issue:
numba's numpy.where doesn't accept scalar inputs for arguments
Example

``` python
import numpy as np
import numba
test_func = lambda x, y: np.where(x<y, x, y)
test_func = numba.jit(nopython=True)(test_func)
test_func(1, 3)
```

``` python
Traceback (most recent call last):
  File "/home/rotis/anaconda/envs/calphad/lib/python2.7/site-packages/numba/typeinfer.py", line 111, in propagate
    constraint(typeinfer)
  File "/home/rotis/anaconda/envs/calphad/lib/python2.7/site-packages/numba/typeinfer.py", line 270, in __call__
    self.resolve(typeinfer, typevars, fnty)
  File "/home/rotis/anaconda/envs/calphad/lib/python2.7/site-packages/numba/typeinfer.py", line 297, in resolve
    sig = context.resolve_function_type(fnty, pos_args, kw_args)
  File "/home/rotis/anaconda/envs/calphad/lib/python2.7/site-packages/numba/typing/context.py", line 113, in resolve_function_type
    return func.get_call_type(self, args, kws)
  File "/home/rotis/anaconda/envs/calphad/lib/python2.7/site-packages/numba/types.py", line 222, in get_call_type
    return self.template(context).apply(args, kws)
  File "/home/rotis/anaconda/envs/calphad/lib/python2.7/site-packages/numba/typing/templates.py", line 155, in apply
    sig = generic(args, kws)
  File "/home/rotis/anaconda/envs/calphad/lib/python2.7/site-packages/numba/typing/npydecl.py", line 747, in generic
    if (cond.ndim == x.ndim == y.ndim and
AttributeError: 'Boolean' object has no attribute 'ndim'
```

Attempts to coerce the scalars into numpy objects (thus with an `ndim` attribute) inside a nopython function have mostly failed.



Discussion:
Replacing `numpy.where` with this snippet seems to work okay as a workaround.

``` python
@numba.vectorize
def nbwhere(condition, x, y):
    if condition:
        return x
    else:
        return y
```


Resolving patch:
diff --git a/numba/targets/arrayobj.py b/numba/targets/arrayobj.py
--- a/numba/targets/arrayobj.py
+++ b/numba/targets/arrayobj.py
@@ -1767,10 +1767,11 @@ def array_nonzero(context, builder, sig, args):
     tup = context.make_tuple(builder, sig.return_type, outs)
     return impl_ret_new_ref(context, builder, sig.return_type, tup)
 
-@builtin
-@implement(numpy.where, types.Kind(types.Array),
-           types.Kind(types.Array), types.Kind(types.Array))
+
 def array_where(context, builder, sig, args):
+    """
+    np.where(array, array, array)
+    """
     layouts = set(a.layout for a in sig.args)
     if layouts == set('C'):
         # Faster implementation for C-contiguous arrays
@@ -1801,6 +1802,24 @@ def where_impl(cond, x, y):
     return impl_ret_untracked(context, builder, sig.return_type, res)
 
 
+@builtin
+@implement(numpy.where, types.Any, types.Any, types.Any)
+def any_where(context, builder, sig, args):
+    cond = sig.args[0]
+    if isinstance(cond, types.Array):
+        return array_where(context, builder, sig, args)
+
+    def scalar_where_impl(cond, x, y):
+        """
+        np.where(scalar, scalar, scalar): return a 0-dim array
+        """
+        scal = x if cond else y
+        return numpy.full_like(scal, scal)
+
+    res = context.compile_internal(builder, scalar_where_impl, sig, args)
+    return impl_ret_new_ref(context, builder, sig.return_type, res)
+
+
 #-------------------------------------------------------------------------------
 # Array attributes
 
@@ -2579,9 +2598,12 @@ def _parse_empty_like_args(context, builder, sig, args):
     np.ones_like() call.
     """
     arytype = sig.args[0]
-    ary = make_array(arytype)(context, builder, value=args[0])
-    shapes = cgutils.unpack_tuple(builder, ary.shape, count=arytype.ndim)
-    return sig.return_type, shapes
+    if isinstance(arytype, types.Array):
+        ary = make_array(arytype)(context, builder, value=args[0])
+        shapes = cgutils.unpack_tuple(builder, ary.shape, count=arytype.ndim)
+        return sig.return_type, shapes
+    else:
+        return sig.return_type, ()
 
 
 @builtin
@@ -2593,8 +2615,8 @@ def numpy_empty_nd(context, builder, sig, args):
     return impl_ret_new_ref(context, builder, sig.return_type, ary._getvalue())
 
 @builtin
-@implement(numpy.empty_like, types.Kind(types.Array))
-@implement(numpy.empty_like, types.Kind(types.Array), types.Kind(types.DTypeSpec))
+@implement(numpy.empty_like, types.Any)
+@implement(numpy.empty_like, types.Any, types.Kind(types.DTypeSpec))
 def numpy_empty_like_nd(context, builder, sig, args):
     arrtype, shapes = _parse_empty_like_args(context, builder, sig, args)
     ary = _empty_nd_impl(context, builder, arrtype, shapes)
@@ -2612,8 +2634,8 @@ def numpy_zeros_nd(context, builder, sig, args):
 
 
 @builtin
-@implement(numpy.zeros_like, types.Kind(types.Array))
-@implement(numpy.zeros_like, types.Kind(types.Array), types.Kind(types.DTypeSpec))
+@implement(numpy.zeros_like, types.Any)
+@implement(numpy.zeros_like, types.Any, types.Kind(types.DTypeSpec))
 def numpy_zeros_like_nd(context, builder, sig, args):
     arrtype, shapes = _parse_empty_like_args(context, builder, sig, args)
     ary = _empty_nd_impl(context, builder, arrtype, shapes)
@@ -2650,7 +2672,7 @@ def full(shape, value, dtype):
 
 
     @builtin
-    @implement(numpy.full_like, types.Kind(types.Array), types.Any)
+    @implement(numpy.full_like, types.Any, types.Any)
     def numpy_full_like_nd(context, builder, sig, args):
 
         def full_like(arr, value):
@@ -2664,7 +2686,7 @@ def full_like(arr, value):
 
 
     @builtin
-    @implement(numpy.full_like, types.Kind(types.Array), types.Any, types.Kind(types.DTypeSpec))
+    @implement(numpy.full_like, types.Any, types.Any, types.Kind(types.DTypeSpec))
     def numpy_full_like_nd(context, builder, sig, args):
 
         def full_like(arr, value, dtype):
@@ -2706,7 +2728,7 @@ def ones(shape, dtype):
     return impl_ret_new_ref(context, builder, sig.return_type, res)
 
 @builtin
-@implement(numpy.ones_like, types.Kind(types.Array))
+@implement(numpy.ones_like, types.Any)
 def numpy_ones_like_nd(context, builder, sig, args):
 
     def ones_like(arr):
@@ -2719,7 +2741,7 @@ def ones_like(arr):
     return impl_ret_new_ref(context, builder, sig.return_type, res)
 
 @builtin
-@implement(numpy.ones_like, types.Kind(types.Array), types.Kind(types.DTypeSpec))
+@implement(numpy.ones_like, types.Any, types.Kind(types.DTypeSpec))
 def numpy_ones_like_dtype_nd(context, builder, sig, args):
 
     def ones_like(arr, dtype):
diff --git a/numba/typing/npydecl.py b/numba/typing/npydecl.py
--- a/numba/typing/npydecl.py
+++ b/numba/typing/npydecl.py
@@ -428,14 +428,23 @@ class NdConstructorLike(CallableTemplate):
     """
 
     def generic(self):
-        def typer(arr, dtype=None):
-            if dtype is None:
-                nb_dtype = arr.dtype
-            else:
+        """
+        np.empty_like(array) -> empty array of the same shape and layout
+        np.empty_like(scalar) -> empty 0-d array of the scalar type
+        """
+        def typer(arg, dtype=None):
+            if dtype is not None:
                 nb_dtype = _parse_dtype(dtype)
+            elif isinstance(arg, types.Array):
+                nb_dtype = arg.dtype
+            else:
+                nb_dtype = arg
             if nb_dtype is not None:
-                layout = arr.layout if arr.layout != 'A' else 'C'
-                return arr.copy(dtype=nb_dtype, layout=layout)
+                if isinstance(arg, types.Array):
+                    layout = arg.layout if arg.layout != 'A' else 'C'
+                    return arg.copy(dtype=nb_dtype, layout=layout)
+                else:
+                    return types.Array(nb_dtype, 0, 'C')
 
         return typer
 
@@ -503,13 +512,22 @@ class NdFullLike(CallableTemplate):
         return_new_reference = True
 
         def generic(self):
-            def typer(arr, fill_value, dtype=None):
-                if dtype is None:
-                    nb_dtype = arr.dtype
-                else:
+            """
+            np.full_like(array, val) -> array of the same shape and layout
+            np.full_like(scalar, val) -> 0-d array of the scalar type
+            """
+            def typer(arg, fill_value, dtype=None):
+                if dtype is not None:
                     nb_dtype = _parse_dtype(dtype)
+                elif isinstance(arg, types.Array):
+                    nb_dtype = arg.dtype
+                else:
+                    nb_dtype = arg
                 if nb_dtype is not None:
-                    return arr.copy(dtype=nb_dtype)
+                    if isinstance(arg, types.Array):
+                        return arg.copy(dtype=nb_dtype)
+                    else:
+                        return types.Array(dtype=nb_dtype, ndim=0, layout='C')
 
             return typer
 
@@ -736,21 +754,30 @@ def generic(self, args, kws):
         assert not kws
 
         if len(args) == 1:
-            # np.where(cond) is the same as np.nonzero(cond)
+            # 0-dim arrays return one result array
             ary = args[0]
             ndim = max(ary.ndim, 1)
             retty = types.UniTuple(types.Array(types.intp, 1, 'C'), ndim)
             return signature(retty, ary)
 
         elif len(args) == 3:
+            # NOTE: contrary to Numpy, we only support homogenous arguments
             cond, x, y = args
-            if (cond.ndim == x.ndim == y.ndim and
-                x.dtype == y.dtype):
-                retty = types.Array(x.dtype, x.ndim, x.layout)
-                return signature(retty, *args)
+            if isinstance(cond, types.Array):
+                # array where()
+                if (cond.ndim == x.ndim == y.ndim and
+                    x.dtype == y.dtype):
+                    retty = types.Array(x.dtype, x.ndim, x.layout)
+                    return signature(retty, *args)
+            else:
+                # scalar where()
+                if not isinstance(x, types.Array) and x == y:
+                    retty = types.Array(x, 0, 'C')
+                    return signature(retty, *args)
 
 builtin_global(numpy.where, types.Function(Where))
 
+
 @builtin
 class Sinc(AbstractTemplate):
     key = numpy.sinc


Test patch:
diff --git a/numba/tests/test_array_methods.py b/numba/tests/test_array_methods.py
--- a/numba/tests/test_array_methods.py
+++ b/numba/tests/test_array_methods.py
@@ -510,6 +510,14 @@ def check_arr(arr):
             self.assertEqual(got.dtype, expected.dtype)
             np.testing.assert_array_equal(got, expected)
 
+        def check_scal(scal):
+            x = 4
+            y = 5
+            cres = compile_isolated(pyfunc, (typeof(scal), typeof(x), typeof(y)))
+            expected = pyfunc(scal, x, y)
+            got = cres.entry_point(scal, x, y)
+            self.assertPreciseEqual(got, expected)
+
         arr = np.int16([1, 0, -1, 0])
         check_arr(arr)
         arr = np.bool_([1, 0, 1])
@@ -527,6 +535,9 @@ def check_arr(arr):
             arr = np.array([v]).reshape(())
             check_arr(arr)
 
+        for x in (0, 1, True, False, 2.5, 0j):
+            check_scal(x)
+
 
 class TestArrayComparisons(TestCase):
 
diff --git a/numba/tests/test_dyn_array.py b/numba/tests/test_dyn_array.py
--- a/numba/tests/test_dyn_array.py
+++ b/numba/tests/test_dyn_array.py
@@ -672,6 +672,9 @@ def check_arr(arr):
             if arr.ndim > 0:
                 check_arr(arr[::2])
 
+        # Scalar argument => should produce a 0-d array
+        check_arr(orig[0])
+
 
 class TestNdEmptyLike(ConstructorLikeBaseTest, TestCase):