1

Consider the following script:

from numba import guvectorize, u1, i8
import numpy as np

@guvectorize([(u1[:],i8)], '(n)->()')
def f(x, res):
    res = x.argmax()

x = np.array([1,2,3],dtype=np.uint8)
print(f(x))
print(x.argmax())
print(f(x))

When running it, I get the following:

4382569440205035030
2
2

Why is this happening? Is there a way to get it right?

Rodrigo Vargas
  • 145
  • 1
  • 3

1 Answers1

0

Python doesn't have references, so res = ... is not actually assigning to the output parameter, but instead rebinding the name res. I believe res is pointing to uninitialized memory, which is why your first run gives a seemingly random value.

Numba works around this using the slice syntax ([:]) which does mutate res- you also need to declare the type as an array. A working function is:

@guvectorize([(u1[:], i8[:])], '(n)->()')
def f(x, res):
    res[:] = x.argmax()
chrisb
  • 49,833
  • 8
  • 70
  • 70
  • I had somehow managed to convince myself that the weird behavior had to do with numba in combination with argmax, but indeed something like `res = x[0]` has the same problem. – Rodrigo Vargas Nov 28 '18 at 18:27