TypingError when using numpy.stack() with numba njit

user_3214872 0
Viewed 45

The original issue is connected to using np.linspace with arrays as start and stop parameters, though right now I'm having issues with the workaround I came up with.

Take the following:

from numba import njit
import numpy as np

@njit
def f1():
  start = np.array([0.1, 1.0], np.float32)
  stop = np.array([1.0, 10.0], np.float32)
  return np.linspace(start, stop, 10)

f1()

This will raise an error, because though documented as supporting "only the 3-argument form" of linspace, what they actually mean is "the 3-argument form with scalar values for start and stop".

So I came up with the folloing workaround:

import numpy as np
from numba import njit

@njit
def f2():
  start = np.array([0.1, 1.0], np.float32)
  stop = np.array([1.0, 10.0], np.float32)
  pts_0 = np.linspace(start[0], stop[0], 10).astype(np.float32) # works
  pts_1 = np.linspace(start[1], stop[1], 10).astype(np.float32) # works
  return np.stack([pts_0, pts_1]).T                             # error

which raises this error:

---------------------------------------------------------------------------
TypingError                               Traceback (most recent call last)
c:\Users\X\Desktop\X\data_analysis.ipynb Cell 46' in <cell line: 18>()
     15   pts_1 = np.linspace(start[1], stop[1], 10).astype(np.float32)
     16   return np.stack([pts_0, pts_1]).T
---> 18 r = f2()

File c:\Users\X\miniconda3\envs\X\lib\site-packages\numba\core\dispatcher.py:468, in _DispatcherBase._compile_for_args(self, *args, **kws)
    464         msg = (f"{str(e).rstrip()} \n\nThis error may have been caused "
    465                f"by the following argument(s):\n{args_str}\n")
    466         e.patch_message(msg)
--> 468     error_rewrite(e, 'typing')
    469 except errors.UnsupportedError as e:
    470     # Something unsupported is present in the user code, add help info
    471     error_rewrite(e, 'unsupported_error')

File c:\Users\X\miniconda3\envs\X\lib\site-packages\numba\core\dispatcher.py:409, in _DispatcherBase._compile_for_args.<locals>.error_rewrite(e, issue_type)
    407     raise e
    408 else:
--> 409     raise e.with_traceback(None)

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<function stack at 0x00000186F280CAF0>) found for signature:
 
 >>> stack(list(array(float32, 1d, C))<iv=None>)

Again, according to the documentation, np.stack is supported (no side-commens on this one either).

What am I missing?

numba python
1 Answers

np.stack is supported but it expect a tuple instead of a list so far. Here is a fixed code:

@njit
def f2():
  start = np.array([0.1, 1.0], np.float32)
  stop = np.array([1.0, 10.0], np.float32)
  pts_0 = np.linspace(start[0], stop[0], 10).astype(np.float32) # works
  pts_1 = np.linspace(start[1], stop[1], 10).astype(np.float32) # works
  return np.stack((pts_0, pts_1)).T                             # works

By the way, note that np.stack((pts_0, pts_1)).T is not very efficient since it creates temporary arrays and a non-contiguous view. Since the purpose of using Numba is to speed up codes, consider using basic loops that should be faster here. The same thing applies for astype(np.float32): a loop can cast the values in-place. Memory and allocations are expensive and this is often what make Numpy slower (also the lack of specific-purpose functions). Such things will be slower in the future (for more information, consider reading more about the "memory wall") so one need to avoid them.

Here is a significantly faster version with basic loops:

@njit
def f2():
    start1, start2 = np.float32(0.1), np.float32(1.0)
    stop1, stop2 = np.float32(1.0), np.float32(10.0)
    steps = 10
    delta = np.float32(1 / (steps - 1))
    res = np.empty((steps, 2), dtype=np.float32)
    for i in range(steps):
        res[i, 0] = start1 + (stop1 - start1) * (delta * i)
        res[i, 1] = start2 + (stop2 - start2) * (delta * i)
    return res

Note that results can be slightly different due to 32-bit FP rounding.

user_129395570
Related