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?
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?
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.