[Python] Calling a numba ufunc from a numba ufunc with an out= argument

(This is a massively stripped down version of the problem. I'm looking for how to solve the compilation error below, not how I ought to have written this toy example.)

Numba let's me write code like this:

import numpy as np
from numba import vectorize

global DOUBLE

@vectorize("int64(int64)")
def double(x):
return 2 * x

DOUBLE = double

@vectorize("int64(int64, int64)")
def two_x_plus_y(x, y):
global DOUBLE
return DOUBLE(x) + y


I can then easily run any of the following code:

arg1 = np.arange(5, dtype=np.int64)
arg2 = arg1 + 5
scratch = np.zeros_like(arg1)

print(two_x_plus_y(arg1, arg2))

result = double(arg1, out=scratch)
assert result is scratch
print(scratch)

print(two_x_plus_y(arg1, arg2, out=scratch))


and everything works exactly as I expected.

But if I write (the somewhat contrived):

@vectorize("int64(int64, int64)")
def two_x_plus_y_alt(x, y):
global DOUBLE
scratch = np.zeros_like(x)
DOUBLE(x, out=scratch)
return scratch + y


The numba compilation fails. Its unhappiness appears to be from the out=. Is there any workaround for this?

In my actual code, I'm trying to avoid making large numbers of large allocations by having all my calculations use a small set of pre-created scratch arrays.

Continue reading...