TypingError when using numpy.stack() with numba njit

Question

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?

Answer 1

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.