My PR to improve argmax broke parallelization, help needed

Hi,

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. https://dev.azure.com/numba/numba/_build/results?buildId=8489&view=logs&j=f0d6d692-dedb-59c9-30d1-aa1d6b8347b2&t=6ceabc93-1c05-51c3-a86f-8b00af8679b9 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: https://github.com/numba/numba/pull/6928#discussion_r613072267.

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