Python numba can't combine arrays-CodePudding

When trying to jit a function that combines two numpy arrays

import numpy as np
import numba as nb

@nb.njit
def combine(a: nb.float64[:], b: nb.float64[:]):
    return np.array([a, b])

Running the function with float parameters doesn't throw an error, i.e.

>>> combine(1., 2.)
array([1., 2.])

But when I try to combine two arrays, I get

>>> combine(np.array([1., 2.]), np.array([3., 4.]))
Traceback (most recent call last):
  File "c:\Users\Lucas Gruwez\Documents\test.py", line 14, in <module>
    combine(np.array([1., 2.]), np.array([3., 4.]))
  File "C:\Users\Lucas Gruwez\AppData\Local\Programs\Python\Python310\lib\site-packages\numba\core\dispatcher.py", line 468, in _compile_for_args
    error_rewrite(e, 'typing')
  File "C:\Users\Lucas Gruwez\AppData\Local\Programs\Python\Python310\lib\site-packages\numba\core\dispatcher.py", line 409, in error_rewrite
    raise e.with_traceback(None)
numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<built-in function array>) found for signature:

 >>> array(list(array(float64, 1d, C))<iv=None>)

There are 4 candidate implementations:
  - Of which 4 did not match due to:
  Overload in function '_OverloadWrapper._build.<locals>.ol_generated': File: numba\core\overload_glue.py: Line 129.
    With argument(s): '(list(array(float64, 1d, C))<iv=None>)':
   Rejected as the implementation raised a specific error:
     TypingError: array(float64, 1d, C) not allowed in a homogeneous sequence
  raised from C:\Users\Lucas Gruwez\AppData\Local\Programs\Python\Python310\lib\site-packages\numba\core\typing\npydecl.py:488

During: resolving callee type: Function(<built-in function array>)
During: typing of call at c:\Users\Lucas Gruwez\Documents\test.py (12)


File "src\test.py", line 12:
def combine(a: nb.float64[:], b: nb.float64[:]):
    return np.array([a, b])

CodePudding user response：

Using np.array([a, b]) to concatenate two numpy arrays is a shortcut which is not allowed due to the strict typing requirements of numba. Numba expects a homogenous sequence like [0,1,2] and therefore fails.

TypingError: array(float64, 1d, C) not allowed in a homogeneous sequence

Instead, only use np.array() to create new arrays from basic types like float or int inside numba and use np.stack or np.concatenate for combining existing arrays.

For example:

@nb.njit
def combine(a: nb.float64[:], b: nb.float64[:]):
    return np.stack((a, b))

combine(np.array([1., 2.]), np.array([3., 4.]))

array([[1., 2.],
       [3., 4.]])

Or:

@nb.njit
def combine(a: nb.float64[:], b: nb.float64[:]):
    return np.concatenate((a, b))

combine(np.array([1., 2.]), np.array([3., 4.]))

array([1., 2., 3., 4.])

PS Due to the static typing, once compile with np.float[:] arrays as dtype, you should not call the function with np.float64 scalars. Likewise, calling np.stack((1.,2.)) with scalar numbers will fail. There you will have to use np.array([1.,2.]) again.

CodePudding user response：

The following code may be helpful, which can be run on both of two single floats or two 1D arrays with the same shapes. For the two arrays, I recommend using indexing to fill the result array, which will work. isinstance is experimental feature now, and can be replaced/written with another modules, e.g. hassattr or np.isscaler, etc, I didn't check if they could be used with numba njit. It was just one solution to show how we could face such problem.

@nb.njit  # (["float64[::1](float64, float64)", "float64[:, ::1](float64[::1], float64[::1])"])
def combine(a, b):
    if isinstance(a, float):
        return np.array([a, b])
    else:
        arr = np.empty((2, *a.shape))
        arr[0, :] = a
        arr[1, :] = b
        return arr