Vectorize with arbitrary output shape

Dear all,

I am trying to parallelize the following function:

@numba.guvectorize([(numba.float64[:,:], numba.int32, numba.int32, numba.float64[:,:])],
                   '(n,m),(),()->(p,q)')
def my_func(array, int1, int2, res):
    for x in range(int1):
        for y in range(int2):
            # DO MANY HEAVY COMPUTATION THAT REQUIRES THE WHOLE ARRAY AND POSITION X AND Y. Note that array shape (n,2) is independent of result shape (which basically depends on the spatial resolution I want)
            res[y, x] = my_result

My problem is that the resulting function has size (int1, int2) but for some reason if p and q are not the same size as the input it throws an error. Is there a way of having a result with an arbitrary size (after all it is an input of the function so it could dynamically allocate memory)?

Thank you for your help,
Jules

1 Like

A not so elegant workaround might be to pass the output array twice, but also as a “dummy” input so the decorator knows what p and q are?

I’m not sure if there’s a more elegant solution. A guvectorize function also return a new array if the output array is omitted (as long as the type is available from the signature). So that might be why it can’t be guaranteed that the output size is known. I’m also never completely clear whether that functionality is actually supported (by design) from Numba, or whether you should always provide the output array.

import numba
import numpy as np

@numba.guvectorize([(
    numba.float64[:,:], 
    numba.int32, 
    numba.int32, 
    numba.float64[:,:],
    numba.float64[:,:],
)], '(n,m),(),(),(p,q)->(p,q)')
def my_func(array, int1, int2, dummy, res):
    for x in range(int1):
        for y in range(int2):
            res[y, x] = x+y

With some example data, that at least seems to work:

x = np.arange(9).reshape(3,3).astype(np.float64)
out = np.full_like(x, np.nan)

my_func(x, 2, 2, out, out);

Resulting out to be:

array([[ 0.,  1., nan],
       [ 1.,  2., nan],
       [nan, nan, nan]])
2 Likes