Intended Behavior? njit'd np.stack does not seem to work with scalars

Hello!
I have a, potentially edge case, issue with a njited function that utilizes np.stack. A minimum working example is below…

import numpy as np
from numba import njit

def f(x, y):

    x2 = x**2
    y2 = y**2

    result = np.stack((x, y))

    return result

fn = njit(f)

x_array = np.linspace(0., 10., 5)
y_array = np.linspace(-1., 1., 5)

x_scalar = 10.
y_scalar = -0.5

The purely Python implementation returns a 2d array when f(x_array, y_array) is called and a 1d array when calling f(x_scalar, y_scalar) (specifically: [10. -0.5]). The njit’d function works fine for fn(x_array, y_array) but fails for fn(x_scalar, y_scalar) with the error:

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<function stack at 0x000002071A9750D0>) found for signature:
>>> stack(UniTuple(float64 x 2))

I am running numba 0.53.1 and numpy 1.19.2 under python 3.9.4.

Does njit not currently support all of np.stack's features? Thanks for the help!

Hi @oniow

I think you have run into a usecase of np.stack that has not been implemented so far. But as far as I can tell, the docs don’t mention the lack of support for non-array inputs (Supported NumPy features — Numba 0.53.1-py3.7-linux-x86_64.egg documentation).

I think it would be good to create an issue on GitHub for this, since it should probably be working.
I think there should not be any obstacles -save time- to implementing a solution.

I think I have seen similar behaviour before, it simply boils down to the fact that there is always a bit of awkwardness when dealing with scalar (single or collections) in numpy (even in CPython), since they fall out of the ordinary array behaviour.

To solve your problem: You could try and add some x = np.asarray(x) (and equivalent for y) in the beginning of the function. In that way you assure that x and y always behave as arrays, and the overhead should be minimal if you are passing actual arrays into f for x and y

Cheers!

Hi @Hannes, thank you for your response! I will make a github issue. I figured it was not a common use case but, as you say, it would be good if it was implemented or at least documented.

Your solution works for now though.

Thanks again!