error from getitem with optional array

[auderson]

Reporting a bug

• I have tried using the latest released version of Numba (most recent is
  visible in the change log (https://github.com/numba/numba/blob/main/CHANGE_LOG).
• I have included a self contained code sample to reproduce the problem.
  i.e. it's possible to run as 'python bug.py'.

Hi, I found the following codes will fail:

import numpy as np
from numba import njit


@njit("float64[:](float64[:], optional(boolean[:]))")
def f1(x, mask=None):
    if mask is None:
        mask = np.isfinite(x)
    return x[mask]

numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<built-in function getitem>) found for signature:
 
 >>> getitem(array(float64, 1d, A), OptionalType(array(bool, 1d, A)))
 
There are 22 candidate implementations:
  - Of which 20 did not match due to:
  Overload of function 'getitem': File: <numerous>: Line N/A.
    With argument(s): '(array(float64, 1d, A), OptionalType(array(bool, 1d, A)))':
   No match.
  - Of which 2 did not match due to:
  Overload in function 'GetItemBuffer.generic': File: numba/core/typing/arraydecl.py: Line 166.
    With argument(s): '(array(float64, 1d, A), OptionalType(array(bool, 1d, A)))':
   Rejected as the implementation raised a specific error:
     NumbaTypeError: unsupported array index type OptionalType(array(bool, 1d, A)) in [OptionalType(array(bool, 1d, A))]
  raised from /home/auderson/miniconda3/lib/python3.8/site-packages/numba/core/typing/arraydecl.py:72
During: typing of intrinsic-call at <ipython-input-4-70443e7e4add> (5)
File "<ipython-input-4-70443e7e4add>", line 5:
def f1(x, mask=None):
    <source elided>
        mask = np.isfinite(x)
    return x[mask]
    ^

@njit("float64[:](float64[:], optional(int64[:]))")
def f2(x, mask=None):
    if mask is None:
        mask = np.arange(x.size)
    return x[mask]

numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<built-in function getitem>) found for signature:
 
 >>> getitem(array(float64, 1d, A), OptionalType(array(int64, 1d, A)))
 
There are 22 candidate implementations:
    - Of which 20 did not match due to:
    Overload of function 'getitem': File: <numerous>: Line N/A.
      With argument(s): '(array(float64, 1d, A), OptionalType(array(int64, 1d, A)))':
     No match.
    - Of which 2 did not match due to:
    Overload in function 'GetItemBuffer.generic': File: numba/core/typing/arraydecl.py: Line 166.
      With argument(s): '(array(float64, 1d, A), OptionalType(array(int64, 1d, A)))':
     Rejected as the implementation raised a specific error:
       NumbaTypeError: unsupported array index type OptionalType(array(int64, 1d, A)) in [OptionalType(array(int64, 1d, A))]
  raised from /home/auderson/miniconda3/lib/python3.8/site-packages/numba/core/typing/arraydecl.py:72
During: typing of intrinsic-call at <ipython-input-5-0db23a20b7d4> (5)
File "<ipython-input-5-0db23a20b7d4>", line 5:
def f2(x, mask=None):
    <source elided>
        mask = np.arange(x.size)
    return x[mask]
    ^

Looks like OptionalType is not correctly handled here.

[guilhermeleobas]

Hi @auderson, thanks for the report. I can confirm this is a bug in Numba.

import numpy as np
from numba import njit

@njit("float64[:](float64[:], optional(boolean[:]))")
def f1(x, mask=None):
    if mask is None:
        mask = np.isfinite(x)
    return x[mask]

x = np.arange(5)
print(f1(x))

Running with NUMBA_DEBUG_TYPEINFER=1, the variable mask (mask.2) at the return point is still typed as Optional:

--------------------------------SSA IR DUMP: f1---------------------------------
label 0:
    x = arg(0, name=x)                       ['x']
    mask = arg(1, name=mask)                 ['mask']
    $const4.1 = const(NoneType, None)        ['$const4.1']
    $6compare_op.2 = mask is $const4.1       ['$6compare_op.2', '$const4.1', 'mask']
    bool8 = global(bool: <class 'bool'>)     ['bool8']
    $8pred = call bool8($6compare_op.2, func=bool8, args=(Var($6compare_op.2, a.py:7),), kws=(), vararg=None, target=None) ['$6compare_op.2', '$8pred', 'bool8']
    branch $8pred, 10, 20                    ['$8pred']
label 10:
    $10load_global.0 = global(np: <module 'numpy' from '/home/guilhermeleobas/miniconda3/envs/numba/lib/python3.8/site-packages/numpy/__init__.py'>) ['$10load_global.0']
    $12load_method.1 = getattr(value=$10load_global.0, attr=isfinite) ['$10load_global.0', '$12load_method.1']
    mask.1 = call $12load_method.1(x, func=$12load_method.1, args=[Var(x, a.py:7)], kws=(), vararg=None, target=None) ['$12load_method.1', 'mask.1', 'x']
    jump 20                                  []
label 20:
    mask.2 = phi(incoming_values=[Var(mask, a.py:7), Var(mask.1, a.py:8)], incoming_blocks=[0, 10]) ['mask', 'mask.1', 'mask.2']
    $24binary_subscr.2 = getitem(value=x, index=mask.2, fn=<built-in function getitem>) ['$24binary_subscr.2', 'mask.2', 'x']
    $26return_value.3 = cast(value=$24binary_subscr.2) ['$24binary_subscr.2', '$26return_value.3']
    return $26return_value.3                 ['$26return_value.3']

-----------------------------------propagate------------------------------------
---- type variables ----
[$10load_global.0 := Module(<module 'numpy' from '/home/guilhermeleobas/miniconda3/envs/numba/lib/python3.8/site-packages/numpy/__init__.py'>),
 $12load_method.1 := Function(<ufunc 'isfinite'>),
 $24binary_subscr.2 := <undecided>,
 $26return_value.3 := array(float64, 1d, A),
 $6compare_op.2 := bool,
 $8pred := bool,
 $const4.1 := none,
 arg.mask := OptionalType(array(bool, 1d, A)),
 arg.x := array(float64, 1d, A),
 bool8 := Function(<class 'bool'>),
 mask := OptionalType(array(bool, 1d, A)),
 mask.1 := array(bool, 1d, C),
 mask.2 := OptionalType(array(bool, 1d, A)),
 x := array(float64, 1d, A)]

[sklam]

TODOs:

• we need to check if Added support for multidimensional array indexing #8491 and [WIP] Advanced Indexing #3: Added support for multiple multidimensional Indices #8912 will fix this.
• if not, we need to extract the fix in Fix: optional array index #7980 and address part 1 in Fix: optional array index #7980 (comment)