refactor np.where to use overload

[guilhermeleobas]

xref: #8254

• Refactor np.where to use @overload
• Broadcast arrays of different shapes. This removes a restriction in the old code where only arrays with the same shape were supported.
• Add NumPy tests
• Add type check for supported types + tests

[guilhermeleobas]

Thanks for the review, @apmasell

[stuartarchibald]

@guilhermeleobas thanks for the patch. In the interests of making this easier to review, perhaps split the API refactoring and the addition of new functionality into separate patches?

[guilhermeleobas]

@stuartarchibald, this is a new impl. of np.where that works with broadcast arrays. You can ignore the diff as none of the old code was actually used.

[stuartarchibald]

@guilhermeleobas please could you resolve the conflicts against main? Many thanks.

[kc611]

The code looks functionally good and well-written. There are mostly issues with handling of None and a few nitpicks but otherwise it LGTM.

[kc611]

Should this be A format ?

Having a C would mean the implementation would tend to go though _where_fast_inner_impl even though arrays might have different layouts (which may not be C or F) ?

(The resulting array should be C ordered but that is automatically handled since the default layout is C ordered.)

[guilhermeleobas]

I just kept the original behavior, which assigns the layout to C:

numba/numba/core/typing/npydecl.py

Lines 710 to 715 in 3bee7be

	if (cond.ndim == x.ndim == y.ndim):
	if x.layout == y.layout == cond.layout:
	retty = types.Array(retdty, x.ndim, x.layout)
	else:
	retty = types.Array(retdty, x.ndim, 'C')
	return signature(retty, *args)

[kc611]

The support for None objects seem to deviate from what's expected:

import numpy as np
import numba

def foo_np(cond, x, y):
    return np.where(cond, x, y)

@numba.njit
def foo_nb(cond, x, y):
    return np.where(cond, x, y)

cond = np.array([None, 1])
x = np.array([0, 1])
y = np.array([3, 4])

print(foo_np(cond, x, y))
print(foo_nb(cond, x, y)) # Error

cond = None
x = np.array([0, 1])
y = np.array([3, 4])

print(foo_np(cond, x, y))
print(foo_nb(cond, x, y)) # Error

cond = np.array([0, 1])
x = np.array([0, 1])
y = None

print(foo_np(cond, x, y))
print(foo_nb(cond, x, y)) # Error

cond = np.array([0, 1])
x = None
y = None

print(foo_np(cond, x, y))
print(foo_nb(cond, x, y)) # Wrong Results

If this is not supported then it should be caught as a proper error. Otherwise, The behaviour of None in condition seem to be the same as 0 and within arrays it seems to be treated as a normal element.

[guilhermeleobas]

Supporting None inputs will be tricky. I'll try to address this in the following days.

[kc611]

Might be better to use the njit and func_name.py_func API over here?

[guilhermeleobas]

I'm following the convention used in the test file.

[kc611]

We could add some tests involving None ?

[stuartarchibald]

Thanks for the patch, this largely looks good, great to see the NumPy tests passing too. I've left a few comments inline, once resolved should be good to merge.

[stuartarchibald]

I think the use of kwargs with default None is going to potentially cause issues. Consider:

In [5]: np.where([3], None, None)
Out[5]: array([None], dtype=object)

vs.

In [13]: np.array([3]).nonzero()
Out[13]: (array([0]),)

[guilhermeleobas]

I've changed the code to not use None as default value.

[stuartarchibald]

I think this should be:

Suggested change

	if x is None and y is None:
	if is_nonelike(x) and is_nonelike(y):

as this is in the typing domain, however, I think there's further issues, e.g.:

from numba import njit
import numpy as np

@njit
def foo(a, x, y):
    return np.where(a, x, y)

args = (1, None, None)
expected = foo.py_func(*args)
got = foo(*args)
print(expected, type(expected))
print(got, type(got))

produces:

None <class 'numpy.ndarray'>
(array([0]),) <class 'tuple'>

setting args = (np.ones(4), None, None) also does something similarly strange.

[stuartarchibald]

This fails to unify for res in the case of 'unusual' inputs like:
cond, x, y = np.arange(-2, 2, 1), np.zeros((4, 4)), np.ones((4, 4), dtype='<U5'))

[guilhermeleobas]

Is there anything we can do in this case?

[stuartarchibald]

Whilst it's maybe possible to assess the output type based on the inputs and explicitly ban unsupported combinations, I think it's ok to leave it as is, the error message is reasonably informative and working out what's "unsupported" is probably complicated. Do you feel differently?

[stuartarchibald]

I guess there are cases where it's faster the compute then broadcast opposed to broadcast then compute. e.g. where the cond is smaller dimension than x and y. Perhaps leave this opt for now and concentrate on correctness with view of getting this merged!

[guilhermeleobas]

@kc611 @stuartarchibald, would you folks be ok with not supporting None inputs in this patch? And instead raise an error when x or y is None.

The current approach fails with None values, as the expression np.asarray(None) is not supported. Once support for it gets included, np.where would work without any major changes.

[stuartarchibald]

@kc611 @stuartarchibald, would you folks be ok with not supporting None inputs in this patch? And instead raise an error when x or y is None.

The current approach fails with None values, as the expression np.asarray(None) is not supported. Once support for it gets included, np.where would work without any major changes.

@guilhermeleobas I think that would be fine, the existing implementation in Numba doesn't support it either so it's not a regression. Thanks!

[stuartarchibald]

Thanks for the updates @guilhermeleobas, think they address everything in the review. I'm inclined to leave the issue with unifying 'unusual' array types for now unless you feel strongly otherwise. Thanks again!