Repository: numba/numba

Issue:
optional type does not work for arrays in njit
<!--

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).

For more general "how do I do X?" type questions, please speak to us in real
time on https://gitter.im/numba/numba or post to the Numba mailing list
https://groups.google.com/a/continuum.io/forum/#!forum/numba-users.

-->

## Reporting a bug

<!--

Before submitting a bug report please ensure that you can check off these boxes:

-->

- [ x ] I am 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 below a minimal working reproducer (if you are unsure how
 to write one see http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports).

<!--

Please include details of the bug here, including, if applicable, what you
expected to happen!

-->

Description:
`optional(float_)` works for njit but `optional(float_[:])` does not work.

Example:

This works:
```python
>>> from numba import optional, float_, njit
>>> @njit((float_, optional(float_)))
... def f(x, y):
...     if y is None:
...         return x
...     return x + y
... 
>>> f(1.0, None)
1.0
>>> f(1.0, 2.0)
3.0

```

But this fails:
```python
>>> @njit((float_[:], optional(float_[:])))
... def f(x, y):
...     if y is None:
...         return x
...     return x + y
... 
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/lizli/anaconda3/lib/python3.7/site-packages/numba/decorators.py", line 198, in wrapper
    disp.compile(sig)
  File "/Users/lizli/anaconda3/lib/python3.7/site-packages/numba/compiler_lock.py", line 32, in _acquire_compile_lock
    return func(*args, **kwargs)
  File "/Users/lizli/anaconda3/lib/python3.7/site-packages/numba/dispatcher.py", line 655, in compile
    cres = self._compiler.compile(args, return_type)
  File "/Users/lizli/anaconda3/lib/python3.7/site-packages/numba/dispatcher.py", line 82, in compile
    pipeline_class=self.pipeline_class)
  File "/Users/lizli/anaconda3/lib/python3.7/site-packages/numba/compiler.py", line 926, in compile_extra
    return pipeline.compile_extra(func)
  File "/Users/lizli/anaconda3/lib/python3.7/site-packages/numba/compiler.py", line 374, in compile_extra
    return self._compile_bytecode()
  File "/Users/lizli/anaconda3/lib/python3.7/site-packages/numba/compiler.py", line 857, in _compile_bytecode
    return self._compile_core()
  File "/Users/lizli/anaconda3/lib/python3.7/site-packages/numba/compiler.py", line 844, in _compile_core
    res = pm.run(self.status)
  File "/Users/lizli/anaconda3/lib/python3.7/site-packages/numba/compiler_lock.py", line 32, in _acquire_compile_lock
    return func(*args, **kwargs)
  File "/Users/lizli/anaconda3/lib/python3.7/site-packages/numba/compiler.py", line 255, in run
    raise patched_exception
  File "/Users/lizli/anaconda3/lib/python3.7/site-packages/numba/compiler.py", line 246, in run
    stage()
  File "/Users/lizli/anaconda3/lib/python3.7/site-packages/numba/compiler.py", line 717, in stage_nopython_backend
    self._backend(lowerfn, objectmode=False)
  File "/Users/lizli/anaconda3/lib/python3.7/site-packages/numba/compiler.py", line 666, in _backend
    lowered = lowerfn()
  File "/Users/lizli/anaconda3/lib/python3.7/site-packages/numba/compiler.py", line 653, in backend_nopython_mode
    self.metadata)
  File "/Users/lizli/anaconda3/lib/python3.7/site-packages/numba/compiler.py", line 1048, in native_lowering_stage
    lower.lower()
  File "/Users/lizli/anaconda3/lib/python3.7/site-packages/numba/lowering.py", line 178, in lower
    self.lower_normal_function(self.fndesc)
  File "/Users/lizli/anaconda3/lib/python3.7/site-packages/numba/lowering.py", line 219, in lower_normal_function
    entry_block_tail = self.lower_function_body()
  File "/Users/lizli/anaconda3/lib/python3.7/site-packages/numba/lowering.py", line 244, in lower_function_body
    self.lower_block(block)
  File "/Users/lizli/anaconda3/lib/python3.7/site-packages/numba/lowering.py", line 259, in lower_block
    self.lower_inst(inst)
  File "/Users/lizli/anaconda3/lib/python3.7/site-packages/numba/lowering.py", line 308, in lower_inst
    val = self.lower_assign(ty, inst)
  File "/Users/lizli/anaconda3/lib/python3.7/site-packages/numba/lowering.py", line 454, in lower_assign
    return self.lower_expr(ty, value)
  File "/Users/lizli/anaconda3/lib/python3.7/site-packages/numba/lowering.py", line 1045, in lower_expr
    res = self.context.special_ops[expr.op](self, expr)
  File "/Users/lizli/anaconda3/lib/python3.7/site-packages/numba/npyufunc/array_exprs.py", line 393, in _lower_array_expr
    caching=False)
  File "/Users/lizli/anaconda3/lib/python3.7/site-packages/numba/targets/base.py", line 843, in compile_subroutine
    flags=flags)
  File "/Users/lizli/anaconda3/lib/python3.7/site-packages/numba/targets/base.py", line 815, in _compile_subroutine_no_cache
    locals=locals)
  File "/Users/lizli/anaconda3/lib/python3.7/site-packages/numba/compiler.py", line 951, in compile_internal
    return pipeline.compile_extra(func)
  File "/Users/lizli/anaconda3/lib/python3.7/site-packages/numba/compiler.py", line 374, in compile_extra
    return self._compile_bytecode()
  File "/Users/lizli/anaconda3/lib/python3.7/site-packages/numba/compiler.py", line 857, in _compile_bytecode
    return self._compile_core()
  File "/Users/lizli/anaconda3/lib/python3.7/site-packages/numba/compiler.py", line 844, in _compile_core
    res = pm.run(self.status)
  File "/Users/lizli/anaconda3/lib/python3.7/site-packages/numba/compiler_lock.py", line 32, in _acquire_compile_lock
    return func(*args, **kwargs)
  File "/Users/lizli/anaconda3/lib/python3.7/site-packages/numba/compiler.py", line 255, in run
    raise patched_exception
  File "/Users/lizli/anaconda3/lib/python3.7/site-packages/numba/compiler.py", line 246, in run
    stage()
  File "/Users/lizli/anaconda3/lib/python3.7/site-packages/numba/compiler.py", line 485, in stage_nopython_frontend
    self.locals)
  File "/Users/lizli/anaconda3/lib/python3.7/site-packages/numba/compiler.py", line 1029, in type_inference_stage
    infer.propagate()
  File "/Users/lizli/anaconda3/lib/python3.7/site-packages/numba/typeinfer.py", line 860, in propagate
    raise errors[0]
numba.errors.TypingError: Failed in nopython mode pipeline (step: nopython mode backend)
Failed in nopython mode pipeline (step: nopython frontend)
No conversion from array(float32, 1d, C) to float32 for '$0.4', defined at None

File "<stdin>", line 3:
<source missing, REPL/exec in use?>

[1] During: typing of assignment at <stdin> (3)

File "<stdin>", line 3:
<source missing, REPL/exec in use?>

[1] During: lowering "$12.3 = arrayexpr(expr=(<built-in function add>, [Var(x, <stdin> (3)), Var(y, <stdin> (3))]), ty=array(float32, 1d, C))" at <stdin> (5)
```


Discussion:
Thanks for the report. I think there are multiple factors governing this behaviour. First, something that does work:
```python
from numba import optional, njit, types
@njit((types.float64, optional(types.float64)))
def f(x, y):
    return x + y
```
and the reason this works is because both `x` and `y` are scalar values, so there's no need to produce a universal function to handle them.

On the other hand, this does not work:
```python
from numba import optional, njit, types
@njit((types.float64[:], optional(types.float64)))
def f(x, y):
    return x + y
```
and produces:
```
unsupported type for input operand: ?float64

File "issue3972.py", line 4:
def f(x, y):
    return x + y
```
which is correct. This happens because the binop `+` is operating on an array type (`x`) which means it gets rewritten as a NumPy ufunc and compiled as such. As NumPy `ufunc`s cannot handle the notion of "Optional" the compilation fails.

The problem you are seeing, which can be generated from:
```python
from numba import optional, njit, types
@njit((types.float64[:], optional(types.float64[:])))
def f(x, y):
    return x + y
```
is because again the binop gets written as a `ufunc` but one of the operands is an optional array. The error happens because there is a bug in the signature handling when generating the ufunc kernel that lets the optional type through as an array, when it should have been rewritten as a scalar/compilation should have been aborted. That the optional array got through is why there's the seemingly bizarre message `"No conversion from array(float32, 1d, C) to float32 for '$0.4'"`, typing is basically confused!

@stuartarchibald Thanks for the explanation. My use case has a bunch of optional arguments (different additional terms in some optimization problem). Is there a good way to implement that with numba at this moment? 

I used to write a small compiler to generate the python code as string based on the input and let numba compile it. The problem was that numba cannot cache functions compiled on the fly #3501  which makes things slow (requiring compiling each session).

I know this approach works for my vector case:
```python
@njit
def f(x, y):
    if y.size:
        return x + y
    return x
```
It is not pretty. I feel this is just tricking the type checker. There might also be edge cases where `numpy.array([])` is semantically different from `None`.

In any case, do you think either this or #3501 is relatively easy to fix? I can help and try to fix it. #3973 is a bit unsatisfactory because it simply prohibits this common use case. If you think the `blah.size` check will not cause future pains for the type checker, I can also add a small note in the `optional` documentation for the vector case.
> Thanks for the explanation.

No problem, I leave answering the "is there a good way to do this" until later!

Fundamentally Numba has to be able to work out the types of everything in the code, further, there are at present some necessary restrictions placed on what can be handled that trade off code re-use, ease of implementation and performance for the common cases.

This:
```python
from numba import optional, njit, types

@njit((types.float64[:], optional(types.float64[:])))
def f(x, y):
    if y.size:
        return x + y
    return x
```
still won't work. That it works in practice without the signature supplied is possibly due to an alteration in the typing of the caller? If this is a reasonable work-around for your use case then perhaps use this for now.

> In any case, do you think either this or #3501 is relatively easy to fix? I can help and try to fix it.

I do not think it is going to be particularly easy to fix, but you are more than welcome to investigate. The Numba core developers do have a workitem scheduled within the next few releases to in general "fix" caching so it just works. Investigating the caching of functions generated at runtime is planned to be part of that effort.

> #3973 is a bit unsatisfactory because it simply prohibits this common use case.

I cannot comment on the frequency of use of this pattern, it has however not been reported in the issue tracker as far as I can recall. #3973 is there to create an error message that offers good feedback to the user about what cannot be compiled and why. This sort of response is often the first thing done because it prevents further confusion/issues and is quick to implement. Unfortunately in this case there are quite a few difficult things to overcome to make it possible to compile the code presented, all of which stem from a compiler pass that rewrites array expressions as `ufunc`s for performance/reuse purposes. `ufunc`s don't have "value or None" as a concept, this is the root cause of the problem, #3973 is in place to prevent this path being taken. It is possible to switch off the rewrite passes, but they are either all on or all off, (other rewrites include but are not limited to stuff to make `print` work and static get/set items), e.g. this works:
```python
from numba import njit
import numpy as np

@njit(no_rewrites=True)
def f(x, y):
    if y is None:
        return x
    return x + y

@njit
def g(x, y):
    if x[0] > 3:
        z = None
    else:
        z = 7 * np.ones((5,))
    return f(x, z)

print(g(1 * np.ones(5), np.zeros(5)))
print(g(10 * np.ones(5), 2 * np.ones(5)))
print(g.signatures)
print(f.signatures)
```
which gives:
```
[8. 8. 8. 8. 8.]
[10. 10. 10. 10. 10.]
[(array(float64, 1d, C), array(float64, 1d, C))]
[(array(float64, 1d, C), ?array(float64, 1d, C))]
```
I'm not sure of the impact of using this in practice as some of the rewrite passes are rather useful, performance may also take a hit!

> My use case has a bunch of optional arguments (different additional terms in some optimization problem). Is there a good way to implement that with numba at this moment?

If Numba can statically determine the value of a predicate involving `None` it will prune the dead branch, for example:

```python
@njit
def f(x, y):
    if y is None:
        return x
    else:
        return x + y

@njit
def call_f(a):
    return f(a, None)

print(call_f(10))
```
will work, as the `else` branch in `f` will be determined as dead ahead of type inference so there's no attempt made to unify `typeof(x) + None`. This also works for `kwargs` e.g.:

```python
from numba import njit
import numpy as np

n = 5

@njit
def myOptimiser(x, y, opt1=None, opt2=None):
    if opt1 is None:
        if opt2 is None:
            return np.ones((n,)) + y
        else:
            return x + opt2
    else:
        if opt2 is None:
            return x * np.ones((n,)) + opt1
        return x - y

@njit
def call_opt(a, b):
    return myOptimiser(a, b, np.ones((n,)))


print(call_opt(10, 20))
```
Hope this helps?





@stuartarchibald Thanks for the help and the suggestions.

1. The proposed implementation unfortunately does not work for my use case, where the presence of the optional arguments are only known at runtime.
2. The intended code is equivalent to the following de-sugared code:
```python
>>> from numba import njit
>>> import numpy as np
>>> 
>>> @njit('f8[:], f8[:], boolean')
... def _f(x, y, y_exists):
...     if y_exists:
...         return x + y
...     return x
... 
>>> def f(x, y=None):
...     if y is None:
...         return _f(x, np.array([]), False)
...     return _f(x, y, True)
... 
>>> f(np.array([1.0]))
array([1.])
>>> f(np.array([1.0]), np.array([2.0]))
array([3.])
```
Maybe the rewrite pass can do this transformation (this does involve not nonpython code though)? We will need a default value for each type (which should not be too bad, at least for `ndarray`.
3. Almost every [numpy ufunc](https://docs.scipy.org/doc/numpy/reference/ufuncs.html) takes optional ndarray arguments... I am surprised nobody bumped into this so far:-)

I will dig into the source code this weekend and see if there is anything I can do for #3972. Do you have a pointer to a good starting point? Thanks~

Resolving patch:
diff --git a/numba/npyufunc/array_exprs.py b/numba/npyufunc/array_exprs.py
--- a/numba/npyufunc/array_exprs.py
+++ b/numba/npyufunc/array_exprs.py
@@ -379,6 +379,8 @@ def _lower_array_expr(lowerer, expr):
     outer_sig = expr.ty(*(lowerer.typeof(name) for name in expr_args))
     inner_sig_args = []
     for argty in outer_sig.args:
+        if isinstance(argty, types.Optional):
+            argty = argty.type
         if isinstance(argty, types.Array):
             inner_sig_args.append(argty.dtype)
         else:
diff --git a/numba/targets/npyimpl.py b/numba/targets/npyimpl.py
--- a/numba/targets/npyimpl.py
+++ b/numba/targets/npyimpl.py
@@ -161,6 +161,14 @@ def _prepare_argument(ctxt, bld, inp, tyinp, where='input operand'):
     _ScalarHelper or _ArrayHelper) class to handle the argument.
     using the polymorphic interface of the Helper classes, scalar
     and array cases can be handled with the same code"""
+
+    # first un-Optional Optionals
+    if isinstance(tyinp, types.Optional):
+        oty = tyinp
+        tyinp = tyinp.type
+        inp = ctxt.cast(bld, inp, oty, tyinp)
+
+    # then prepare the arg for a concrete instance
     if isinstance(tyinp, types.ArrayCompatible):
         ary     = ctxt.make_array(tyinp)(ctxt, bld, inp)
         shape   = cgutils.unpack_tuple(bld, ary.shape, tyinp.ndim)


Test patch:
diff --git a/numba/tests/test_array_exprs.py b/numba/tests/test_array_exprs.py
--- a/numba/tests/test_array_exprs.py
+++ b/numba/tests/test_array_exprs.py
@@ -6,7 +6,7 @@
 
 from numba import njit, vectorize
 from numba import unittest_support as unittest
-from numba import compiler, typing, typeof, ir, utils
+from numba import compiler, typing, typeof, ir, utils, types
 from numba.compiler import Pipeline, _PipelineManager, Flags
 from numba.targets import cpu
 from .support import MemoryLeakMixin, TestCase
@@ -486,5 +486,133 @@ def test_division_by_zero(self):
         np.testing.assert_array_equal(expect, got)
 
 
+class TestOptionals(MemoryLeakMixin, unittest.TestCase):
+    """ Tests the arrival and correct lowering of Optional types at a arrayexpr
+    derived ufunc, see #3972"""
+
+    def test_optional_scalar_type(self):
+
+        @njit
+        def arr_expr(x, y):
+            return x + y
+
+        @njit
+        def do_call(x, y):
+            if y > 0:
+                z = None
+            else:
+                z = y
+            return arr_expr(x, z)
+
+        args = (np.arange(5), -1.2)
+
+        # check result
+        res = do_call(*args)
+        expected = do_call.py_func(*args)
+        np.testing.assert_allclose(res, expected)
+
+        # check type
+        s = arr_expr.signatures
+        oty = s[0][1]
+        self.assertTrue(isinstance(oty, types.Optional))
+        self.assertTrue(isinstance(oty.type, types.Float))
+
+
+    def test_optional_array_type(self):
+
+        @njit
+        def arr_expr(x, y):
+            return x + y
+
+        @njit
+        def do_call(x, y):
+            if y[0] > 0:
+                z = None
+            else:
+                z = y
+            return arr_expr(x, z)
+
+        args = (np.arange(5), np.arange(5.))
+
+        # check result
+        res = do_call(*args)
+        expected = do_call.py_func(*args)
+        np.testing.assert_allclose(res, expected)
+
+        # check type
+        s = arr_expr.signatures
+        oty = s[0][1]
+        self.assertTrue(isinstance(oty, types.Optional))
+        self.assertTrue(isinstance(oty.type, types.Array))
+        self.assertTrue(isinstance(oty.type.dtype, types.Float))
+
+
+class TestOptionalsExceptions(MemoryLeakMixin, unittest.TestCase):
+    # same as above, but the Optional resolves to None and TypeError's
+
+    def test_optional_scalar_type_exception_on_none(self):
+
+        self.disable_leak_check()
+
+        @njit
+        def arr_expr(x, y):
+            return x + y
+
+        @njit
+        def do_call(x, y):
+            if y > 0:
+                z = None
+            else:
+                z = y
+            return arr_expr(x, z)
+
+        args = (np.arange(5), 1.0)
+
+        # check result
+        with self.assertRaises(TypeError) as raises:
+            do_call(*args)
+
+        self.assertIn("expected float64, got None", str(raises.exception))
+
+        # check type
+        s = arr_expr.signatures
+        oty = s[0][1]
+        self.assertTrue(isinstance(oty, types.Optional))
+        self.assertTrue(isinstance(oty.type, types.Float))
+
+    def test_optional_array_type_exception_on_none(self):
+
+        self.disable_leak_check()
+
+        @njit
+        def arr_expr(x, y):
+            return x + y
+
+        @njit
+        def do_call(x, y):
+            if y[0] > 0:
+                z = None
+            else:
+                z = y
+            return arr_expr(x, z)
+
+        args = (np.arange(5), np.arange(1., 5.))
+
+        # check result
+        with self.assertRaises(TypeError) as raises:
+            do_call(*args)
+
+        excstr = str(raises.exception)
+        self.assertIn("expected array(float64,", excstr)
+        self.assertIn("got None", excstr)
+
+        # check type
+        s = arr_expr.signatures
+        oty = s[0][1]
+        self.assertTrue(isinstance(oty, types.Optional))
+        self.assertTrue(isinstance(oty.type, types.Array))
+        self.assertTrue(isinstance(oty.type.dtype, types.Float))
+
+
 if __name__ == "__main__":
     unittest.main()