Broadcasting/Shape Exception in Numba guvectorize Signature

I'm trying to use the @guvectorize decorator that will return the axis=0 means of an array from vectors of starts and stops, so for example, if:

values = array([[82, 3],                  [76, 7],                  [23, 8],                  [46, 5],                  [38, 6]])  starts = array([ 0,  1,  1,  1,  4])  stops  = array([ 1,  4,  4,  4,  5])  

My function would return:

array([[82.        ,  3.        ],         [48.33333333,  6.66666667],         [48.33333333,  6.66666667],         [48.33333333,  6.66666667],         [38.        ,  6.        ]])  

Following the simple example in the numba documentation, this is what I came up with (function operates over the vector of one column):

@guvectorize([(int64[:], int64[:], int64[:], float64[:])], '(n),(n),(n)->(n)', nopython=True)  def variable_window_avg(values, starts, stops, means):      for i in range(values.shape[0]):          means[i] = np.nanmean(values[starts[i]:stops[i]])  

This works great when values is a vector, but does not scale up when values is an ndarray as desired:

---------------------------------------------------------------------------  ValueError                                Traceback (most recent call last)  <ipython-input-128-8a213e6ace9c> in <module>  ----> 1 variable_window_avg(arr[:,:2], arr[:,1], arr[:,2])    ValueError: variable_window_avg: Input operand 1 has a mismatch in its core dimension 0, with gufunc signature (n),(n),(n)->(n) (size 100 is different from 2)  

The error suggests the shapes are not broadcasting together correctly via the Numpy generalized universal functions API, and indeed

res = variable_window_avg(values.T, starts, stops)    res = res.T  

works. But modifying the input and output layouts signature would seem the better solution to avoid two transposes. I also don't understand why the example in the numba documentation but my code fails.

https://stackoverflow.com/questions/67029825/broadcasting-shape-exception-in-numba-guvectorize-signature April 10, 2021 at 08:45AM