Repository: numba/numba

Issue:
numpy.argmin(arr, axis=int64), where (arr.ndim > 1), returns passes through the dtype of arr
When operating on a 1D array within a JIT-ed function, `numpy.argmin` always returns an array of type int64, regardless of whether or not you specify `axis=0`. However, when operating on a higher-dimension array, the dtype of the array is passed through to the return value. Here is a simple demonstration script:

```python
import numpy as np
from numba import njit


def test():
    one = np.argmin(np.arange(1024) < 512, axis=0)
    two_zero = np.argmin(np.arange(9).reshape((3, 3)).T < 5, axis=0)
    two_one = np.argmin(np.arange(9).reshape((3, 3)) < 4, axis=1)
    x = (np.arange(1023, -1, -1).reshape((512, 2)) // 4).astype(np.uint8)
    another_test = np.argmin(x, axis=0)
    return one, two_zero, two_one, another_test


if __name__ == '__main__':
    print(test())  # (512, array([0, 2, 0], dtype=int64), array([0, 1, 0], dtype=int64), array([510, 510], dtype=int64))
    print(njit(test)())  # (512, array([False,  True, False]), array([False,  True, False]), array([254, 254], dtype=uint8))
```

This can be avoided by casting the array with `.astype('int64')`, but its rather strange that you would need to do this.


Discussion:
Bug confirmed. `np.argmin` return type should not be based on the input array type: https://github.com/numba/numba/blob/290bcaa63804fb99b1fe238b0f8fe1caf7686899/numba/np/arraymath.py#L798



Resolving patch:
diff --git a/numba/np/arraymath.py b/numba/np/arraymath.py
--- a/numba/np/arraymath.py
+++ b/numba/np/arraymath.py
@@ -795,7 +795,7 @@ def build_argmax_or_argmin_with_axis_impl(arr, axis, flatten_impl):
     array, return the implementation function.
     """
     check_is_integer(axis, "axis")
-    retty = arr.dtype
+    retty = types.intp
 
     tuple_buffer = tuple(range(arr.ndim))
 


Test patch:
diff --git a/numba/tests/test_array_reductions.py b/numba/tests/test_array_reductions.py
--- a/numba/tests/test_array_reductions.py
+++ b/numba/tests/test_array_reductions.py
@@ -1021,6 +1021,16 @@ def argmax(arr):
         self.assertPreciseEqual(argmax(arr2d),
                                 jit(nopython=True)(argmax)(arr2d))
 
+    def test_argmax_return_type(self):
+        # See issue #7853, return type should be intp not based on input type
+        arr2d = np.arange(6, dtype=np.uint8).reshape(2, 3)
+
+        def argmax(arr):
+            return arr2d.argmax(axis=0)
+
+        self.assertPreciseEqual(argmax(arr2d),
+                                jit(nopython=True)(argmax)(arr2d))
+
     def test_argmin_axis_1d_2d_4d(self):
         arr1d = np.array([0, 20, 3, 4])
         arr2d = np.arange(6).reshape(2, 3)
@@ -1083,6 +1093,16 @@ def argmin(arr):
         self.assertPreciseEqual(argmin(arr2d),
                                 jit(nopython=True)(argmin)(arr2d))
 
+    def test_argmin_return_type(self):
+        # See issue #7853, return type should be intp not based on input type
+        arr2d = np.arange(6, dtype=np.uint8).reshape(2, 3)
+
+        def argmin(arr):
+            return arr2d.argmin(axis=0)
+
+        self.assertPreciseEqual(argmin(arr2d),
+                                jit(nopython=True)(argmin)(arr2d))
+
     @classmethod
     def install_generated_tests(cls):
         # These form a testing product where each of the combinations are tested