Exclude arguments from vectorisation

Numba Support: How do I do ...?
1

I have a function:

func(x, arr)

I wish to use numba to vectorize over x, but exclude arr. This is done with numpy by using np.vectorize(func, exclude=[1]). Can this be done with numba?

2

You can do it using guvectorize, but it will require you to provide some information about the dimensions.

Here is a silly example where arr is 2D, as specified in the signature, but x can be any shape and the function will be applied for one element of x at the time.

from numba import guvectorize
import numpy as np

@guvectorize("(),(n,m)->()")
def myfunc(x, arr, out):
    
    for idx in np.ndindex(arr.shape):
        x += arr[idx]
        
    out[:] = x

x = np.random.rand(3,3,3)
arr = np.arange(9).reshape(3,3)
out = np.empty_like(x)

myfunc(x, arr, out)

assert np.allclose(out - arr.sum(), x)

In the above case you can of course generate the signature on-the-fly by examining the shape of arr, in case it’s not constant.