My PR to improve argmax broke parallelization, help needed


I have submitted a PR to add support for axis keyword argument to np.argmax: [WIP] Add preliminary support for axis to argmax by itamarst · Pull Request #6928 · numba/numba · GitHub

The functionality appears to work, but I’ve… broken parallelization, and I have no idea how re-enable it. shows the failing test.

I’m not really sure how to reenable parallelization (was switching to @overload a mistake? I couldn’t get @lower_builtin to work), so any advice would be welcome.

1 Like

Hi @itamarst

Thanks for working on implementing this. I’ve taken a look at the code and commented on the PR with what I think the issue is:

(was switching to @overload a mistake? I couldn’t get @lower_builtin to work),

Using @overload is in general the preferred route for implementing functionality such as NumPy functions in Numba. There’s a guide to help here: A guide to using @overload — Numba 0+untagged.4124.gd4460fe.dirty documentation

Hope this helps.