Repository: numba/numba

Issue:
Type inferring bug
<!--

Thanks for opening an issue! To help the Numba team handle your information
efficiently, please first ensure that there is no other issue present that
already describes the issue you have
(search at https://github.com/numba/numba/issues?&q=is%3Aissue).

-->

## Reporting a bug

<!--

Before submitting a bug report please ensure that you can check off these boxes:

-->

- [x] I have tried using the latest released version of Numba (most recent is
 visible in the change log (https://github.com/numba/numba/blob/master/CHANGE_LOG).
- [x] I have included a self contained code sample to reproduce the problem.
  i.e. it's possible to run as 'python bug.py'.

<!--

Please include details of the bug here, including, if applicable, what you
expected to happen!

-->

It seems I found a type inferring issue. Here is an example to reproduce it:

```
@njit
def inner(x, x_=None):
    if x_ is not None:
        y = x_
    else:
        y = x + 1
    return np.add(y, 1)

@njit
def test(a):
    if a > 0:
        b = np.arange(5.)
    else:
        b = None
    return inner(np.arange(5.), b)

>>> test(1)
```
this produces the following error:
```
Traceback (most recent call last):
  File "/home/auderson/miniconda3/lib/python3.8/site-packages/IPython/core/interactiveshell.py", line 3437, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-203-81ccd7593ad9>", line 1, in <module>
    test(1)
  File "/home/auderson/miniconda3/lib/python3.8/site-packages/numba/core/dispatcher.py", line 420, in _compile_for_args
    error_rewrite(e, 'typing')
  File "/home/auderson/miniconda3/lib/python3.8/site-packages/numba/core/dispatcher.py", line 361, in error_rewrite
    raise e.with_traceback(None)
numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<ufunc 'add'>) found for signature:
 
 >>> add(OptionalType(array(float64, 1d, C)) i.e. the type 'array(float64, 1d, C) or None', Literal[int](1))
 
There are 2 candidate implementations:
      - Of which 1 did not match due to:
      Overload in function 'Numpy_rules_ufunc.generic': File: numba/core/typing/npydecl.py: Line 96.
        With argument(s): '(OptionalType(array(float64, 1d, C)) i.e. the type 'array(float64, 1d, C) or None', int64)':
       Rejected as the implementation raised a specific error:
         TypingError: can't resolve ufunc add for types (OptionalType(array(float64, 1d, C)), int64)
  raised from /home/auderson/miniconda3/lib/python3.8/site-packages/numba/core/typing/npydecl.py:102
      - Of which 1 did not match due to:
      Overload in function 'Numpy_rules_ufunc.generic': File: numba/core/typing/npydecl.py: Line 96.
        With argument(s): '(OptionalType(array(float64, 1d, C)) i.e. the type 'array(float64, 1d, C) or None', Literal[int](1))':
       Rejected as the implementation raised a specific error:
         TypingError: can't resolve ufunc add for types (OptionalType(array(float64, 1d, C)), Literal[int](1))
  raised from /home/auderson/miniconda3/lib/python3.8/site-packages/numba/core/typing/npydecl.py:102
During: resolving callee type: Function(<ufunc 'add'>)
During: typing of call at <ipython-input-202-1466296e42c7> (7)
File "<ipython-input-202-1466296e42c7>", line 7:
def inner(x, x_=None):
    <source elided>
        y = x + 1
    return np.add(y, 1)
    ^
During: resolving callee type: type(CPUDispatcher(<function inner at 0x7f29bc1adaf0>))
During: typing of call at <ipython-input-202-1466296e42c7> (14)
During: resolving callee type: type(CPUDispatcher(<function inner at 0x7f29bc1adaf0>))
During: typing of call at <ipython-input-202-1466296e42c7> (14)
File "<ipython-input-202-1466296e42c7>", line 14:
def test(a):
    <source elided>
        b = None
    return inner(np.arange(5.), b)
    ^
```

as you can see, variable `b` from `test` is `OptionalType(array(float64, 1d, C))`, which caused the error. 
but in `inner`, there is a check for `None`, so `y` ‘s type should not be `OptionalType(array(float64, 1d, C))` any more.

if change `np.add` to `+`, it will cause a `LoweringError`

if it's not easy to change how type inferring works, do I have other ways to deal with `OptionalType(array(float64, 1d, C))` ? i.e. how can I make it a finite array type  inside the `inner` function (I do need the argument to be **optional**)?

my numba version is '0.53.1'


Discussion:
Thanks @auderson for reporting this issue and providing a reproducer. 
I was able to reproduce it as well.
I believe the problem is in the `if-else` statement in `test`  and how the type for `b` is resolved at compile time.
If you replaced your if-else structure to be like the following, it will work
```python
@njit
def test(a):
    if a > 0:
        b = np.arange(5.)
        return inner(np.arange(5.), b)
    else:
        return inner(np.arange(5.))
```
@hadia206 

Yes, this will work, as in this version, `b` is not a `OptionalType(array(float64, 1d, C))`. However if I use `if-else` in all the cases to make `b`'s type not optional, my code will be very complicated and redundant.

But luckily, later yesterday I found a way to get around:
just wrap the `y` in `inner` with `np.asarray`:

```
@njit
def inner(x, x_=None):
    if x_ is not None:
        y = x_
    else:
        y = x + 1
    return np.add(np.asarray(y), 1)
```
And what interests me here is that can we enable the compiler to infer `y`'s type more precisely, given I've known `x_` is already not `None`, so that it's not of `OptionalType(array(float64, 1d, C))` but `array` type?

Resolving patch:
diff --git a/numba/core/typing/npydecl.py b/numba/core/typing/npydecl.py
--- a/numba/core/typing/npydecl.py
+++ b/numba/core/typing/npydecl.py
@@ -95,6 +95,9 @@ def ufunc(self):
         return self.key
 
     def generic(self, args, kws):
+        # First, strip optional types, ufunc loops are typed on concrete types
+        args = [x.type if isinstance(x, types.Optional) else x for x in args]
+
         ufunc = self.ufunc
         base_types, explicit_outputs, ndims, layout = self._handle_inputs(
             ufunc, args, kws)
diff --git a/numba/np/npyimpl.py b/numba/np/npyimpl.py
--- a/numba/np/npyimpl.py
+++ b/numba/np/npyimpl.py
@@ -232,6 +232,10 @@ def _build_array(context, builder, array_ty, input_types, inputs):
     given the target context, builder, output array type, and a list of
     _ArrayHelper instances.
     """
+    # First, strip optional types, ufunc loops are typed on concrete types
+    input_types = [x.type if isinstance(x, types.Optional) else x
+                   for x in input_types]
+
     intp_ty = context.get_value_type(types.intp)
     def make_intp_const(val):
         return context.get_constant(types.intp, val)


Test patch:
diff --git a/numba/tests/npyufunc/test_ufunc.py b/numba/tests/npyufunc/test_ufunc.py
--- a/numba/tests/npyufunc/test_ufunc.py
+++ b/numba/tests/npyufunc/test_ufunc.py
@@ -1,6 +1,6 @@
 import numpy as np
 
-from numba import float32, jit
+from numba import float32, jit, njit
 from numba.np.ufunc import Vectorize
 from numba.core.errors import TypingError
 from numba.tests.support import TestCase
@@ -117,6 +117,29 @@ def test(x):
 
             self.assertIn(msg, str(raises.exception))
 
+    def test_optional_type_handling(self):
+        # Tests ufunc compilation with Optional type
+
+        @njit
+        def inner(x, y):
+            if y > 2:
+                z = None
+            else:
+                z = np.ones(4)
+            return np.add(x, z)
+
+        # This causes `z` to be np.ones(4) at runtime, success
+        self.assertPreciseEqual(inner(np.arange(4), 1),
+                                np.arange(1, 5).astype(np.float64))
+
+        with self.assertRaises(TypeError) as raises:
+            # This causes `z` to be None at runtime, TypeError raised on the
+            # type cast of the Optional.
+            inner(np.arange(4), 3)
+
+        msg = "expected array(float64, 1d, C), got None"
+        self.assertIn(msg, str(raises.exception))
+
 
 if __name__ == '__main__':
     unittest.main()