np.empty support list type in argument

[ZeppLu]

import numpy as np
import numba as nb

tupleShape = nb.njit(lambda N: np.zeros((N,N)))
listShape = nb.njit(lambda N: np.zeros([N,N]))

tupleShape(10)  # fine
listShape(10)  # TypingError: Failed at nopython (nopython frontend)
               # Invalid usage of Function(<built-in function zeros>) with parameters (list(int64))
               #  * parameterized
               # In definition 0:
               #     All templates rejected
               # [1] During: resolving callee type: Function(<built-in function zeros>)
               # [2] During: typing of call at <ipython-input-3-94d204f55617> (5)

I know it should be a small issue, anyone comes across it can resolve by simply altering his code a little bit. But for those used to MATLAB, the small issue is just annoying.

[stuartarchibald]

Thanks for the report. I think the problem here is that np.zeros(), and in fact all the routines which are backed by np.empty for allocation, do not support list type as an argument.

[stuartarchibald]

note: Edited title to reflect problem.

[umangv]

I'm running into this problem because np.empty and np.zeros accept neither lists nor numpy arrays.

For example,

@njit
def zeros(d):
    return np.zeros(np.full(d, 2))

raises a TypingError

[stuartarchibald]

@umangv thanks for the report. As noted above, lists are not supported for anything that needs np.empty calls for allocation. Lists are hard because at compile time Numba only knows that it is a List type, it doesn't know how long it is, and the length determines the type of the returned allocation, e.g. np.empty([1, 2, 3]) is a 3D array. Without the type information being available compilation cannot be performed. In the case of array input, the same problem is present, the contents of the array at runtime determines the size of the allocation, whereas at compile time all is know is that e.g. np.full(d, 2) is going to return a 1D array. For this sort of thing to work more generally Numba will need to have support for value based dispatch #2949 and even then supporting this for more complex types will likely be hard.

[umangv]

Interesting, so you need to know the dimension of the desired array at compile time but not necessarily the shape/size. That was the distinction that I was missing.

[gdonval]

Was about to open a bug report (instead of a feature request) about this: numpy does support it and this says that Any will do...

Anyway, it is possible to work around the problem in some (very) specific cases since #2691 without too much trickery or having to get back to the interpreter:

import numba as nb
import numpy as np 
from numba.unsafe.ndarray import to_fixed_tuple

def make_allocator(ndims):
    _ndims = int(ndims)
    def _allocator_wrapper(shape_arr):
        _shape = to_fixed_tuple(shape_arr, _ndims)
        allocated = np.zeros(_shape, dtype=np.float64)
        return allocated
    return nb.njit(_allocator_wrapper)

# OK with arrays:
allocate_2D_array = make_allocator(2)
allocate_3D_array = make_allocator(3)

It's not really solving the problem but in specific cases where the number of dimensions is known from the start, everything can be JITed like that without exiting Numba.