Numba compiled function cannot iterate over tuple of numpy arrays with different dimensions

Question

I am trying to modify some numpy arrays with different dimensions, passed to numba function with a tuple. The function looks like this:

from numba import njit

def numbatest(tuple1,tuple2):
    for i in range(2):
        tuple1[i][0]=tuple2[i][2]
    return tuple1,tuple2

numbatest_njit=njit(numbatest)

When I pass it the the tuples of two numpy arrays (1D and 2D):

a=np.empty(10)
b=np.empty([10,3])
c=np.empty(10)
d=np.empty([10,3])
A=(a,b)
B=(c,d)
C,D=numbatest_njit(A,B)

I get the following error

TypingError                               Traceback (most recent call last)
<ipython-input-179-f6c1f66607ba> in <module>
      6 A=(a,b)
      7 B=(c,d)
----> 8 C,D=numbatest_njit(A,B)
      9 

~/.local/lib/python3.6/site-packages/numba/core/dispatcher.py in _compile_for_args(self, *args, **kws)
    418                 e.patch_message(msg)
    419 
--> 420             error_rewrite(e, 'typing')
    421         except errors.UnsupportedError as e:
    422             # Something unsupported is present in the user code, add help info

~/.local/lib/python3.6/site-packages/numba/core/dispatcher.py in error_rewrite(e, issue_type)
    359                 raise e
    360             else:
--> 361                 raise e.with_traceback(None)
    362 
    363         argtypes = []

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<built-in function getitem>) found for signature:
 
 >>> getitem(Tuple(array(float64, 1d, C), array(float64, 2d, C)), int64)
 
There are 22 candidate implementations:
      - Of which 22 did not match due to:
      Overload of function 'getitem': File: <numerous>: Line N/A.
        With argument(s): '(Tuple(array(float64, 1d, C), array(float64, 2d, C)), int64)':
       No match.

During: typing of intrinsic-call at <ipython-input-178-e1c23213087e> (6)

File "<ipython-input-178-e1c23213087e>", line 6:
def numbatest(tuple1,tuple2):
    <source elided>
    for i in range(2):
        tuple1[i][0]=tuple2[i][2]
        ^

Interestingly, the function works if I rewrite it so that I explicitly call different tuples:

def numbatest(tuple1,tuple2):
    i=0
    tuple1[i][0]=tuple2[i][2]
    i=1
    tuple1[i][0]=tuple2[i][2]
    return tuple1,tuple2

The first function works if it is called directly (without compling it with numba). If all of the numpy arrays are 1D arrays, then it also works with numba.

Does anyone know what could be the problem here? And how to avoid it? How would you pass the numpy arrays of different sizes to functions instead? Thank you very much in advance!

Answer 1

Numba supports different sets of operations for homogeneous and heterogeneous tuples. In particular, when it comes to tuple indexing, numba only supports indexing heterogeneous tuples if the index is a constant known at compile-time.

This explains the difference in behaviour when the tuples contain 1D arrays (homogeneous tuples) vs a mixture of 1D and 2D arrays (heterogeneous tuples). It also explains why the second version of the function works: numba can tell that i = 0 and i = 1 are compile-time constants, but not smart enough to know that range(2) could be unrolled to the same thing.