Add support for np.diag_indices() and np.diag_indices_from()

docs/source/reference/numpysupported.rst

@@ -467,6 +467,8 @@ The following top-level functions are supported:
 
 * :func:`numpy.delete` (only the 2 first arguments)
 * :func:`numpy.diag`
+* :func:`numpy.diag_indices` (returns a NumPy array instead of a tuple)
+* :func:`numpy.diag_indices_from` (returns a NumPy array instead of a tuple)
 * :func:`numpy.digitize`
 * :func:`numpy.dstack`
 * :func:`numpy.dtype` (only the first argument)

docs/upcoming_changes/9115.np_support.rst

@@ -0,0 +1,4 @@
+
+Add np.diag_indices() and np.diag_indices_from()
+================================================
+Support for np.diag_indices() and np.diag_indices_from() is added.

numba/np/arrayobj.py

@@ -4523,6 +4523,45 @@ def impl(N, M=None, k=0, dtype=float):
     return impl
 
 
+@overload(np.diag_indices)
+def np_diag_indices(n, ndim=2):
+    if not isinstance(n, types.Integer):
+        msg = 'The argument "n" must be an integer'
+        raise errors.TypingError(msg)
+
+    if not isinstance(ndim, (int, types.Integer)):
+        msg = 'The argument "ndim" must be an integer'
+        raise errors.TypingError(msg)
+
+    def impl(n, ndim=2):
+        res = np.arange(n * ndim)
+        for i in range(n * ndim):
+            res[i] = res[i] % n
+        return res.reshape((ndim, n))
+
+    return impl
+
+
+@overload(np.diag_indices_from)
+def np_diag_indices_from(arr):
+    if not isinstance(arr, types.Array):
+        msg = 'The argument "arr" must be an array'
+        raise errors.TypingError(msg)
+
+    def impl(arr):
+        if not arr.ndim >= 2:
+            raise ValueError("Input array must be at least 2-d")
+        # For more than d=2, the strided formula is only valid for arrays with
+        # all dimensions equal, so we check first.
+        s = np.asarray(arr.shape)
+        if not np.all(np.diff(s) == 0):
+            raise ValueError("All dimensions of input must be of equal length")
+
+        return np.diag_indices(arr.shape[0], arr.ndim)
+
+    return impl
+
+
 @overload(np.diag)
 def impl_np_diag(v, k=0):
     if not type_can_asarray(v):

numba/tests/test_np_functions.py

@@ -432,6 +432,18 @@ def nan_to_num(X, copy=True, nan=0.0):
     return np.nan_to_num(X, copy=copy, nan=nan)
 
 
+def diag_indices1(n):
+    return np.diag_indices(n)
+
+
+def diag_indices2(n, ndim):
+    return np.diag_indices(n, ndim)
+
+
+def diag_indices_from(arr):
+    return np.diag_indices_from(arr)
+
+
 class TestNPFunctions(MemoryLeakMixin, TestCase):
     """
     Tests for various Numpy functions.
@@ -5547,6 +5559,103 @@ def test_nan_to_num_invalid_argument(self):
         self.assertIn("The first argument must be a scalar or an array-like",
                       str(raises.exception))
 
+    def test_diag_indices_basic(self):
+        pyfunc1 = diag_indices1
+        cfunc1 = njit(pyfunc1)
+        pyfunc2 = diag_indices2
+        cfunc2 = njit(pyfunc2)
+
+        def inputs():
+            # Based on https://github.com/numpy/numpy/blob/...
+            # 09ee7daa9113ecb503552913d76185d0ad3f0f31/numpy/lib/tests/...
+            # test_index_tricks.py#L483-L508
+            yield 4, 2
+            yield 2, 3
+
+        for n, ndim in inputs():
+            self.assertPreciseEqual(np.asarray(pyfunc1(n)), cfunc1(n))
+            self.assertPreciseEqual(np.asarray(pyfunc2(n, ndim)),
+                                    cfunc2(n, ndim))
+
+    def test_diag_indices1_exception(self):
+        pyfunc = diag_indices1
+        cfunc = njit(pyfunc)
+
+        self.disable_leak_check()
+
+        with self.assertRaises(TypingError) as raises:
+            cfunc("abc")
+        self.assertIn('The argument "n" must be an integer',
+                      str(raises.exception))
+
+        with self.assertRaises(TypingError) as raises:
+            cfunc(3.0)
+        self.assertIn('The argument "n" must be an integer',
+                      str(raises.exception))
+
+    def test_diag_indices2_exception(self):
+        pyfunc = diag_indices2
+        cfunc = njit(pyfunc)
+
+        self.disable_leak_check()
+
+        with self.assertRaises(TypingError) as raises:
+            cfunc("abc", 2)
+        self.assertIn('The argument "n" must be an integer',
+                      str(raises.exception))
+
+        with self.assertRaises(TypingError) as raises:
+            cfunc(4, "abc")
+        self.assertIn('The argument "ndim" must be an integer',
+                      str(raises.exception))
+
+        with self.assertRaises(TypingError) as raises:
+            cfunc(4, 3.0)
+        self.assertIn('The argument "ndim" must be an integer',
+                      str(raises.exception))
+
+    def test_diag_indices_from_basic(self):
+        pyfunc = diag_indices_from
+        cfunc = njit(pyfunc)
+
+        def inputs():
+            # Taken from https://github.com/numpy/numpy/blob/...
+            # 09ee7daa9113ecb503552913d76185d0ad3f0f31/numpy/lib/tests/...
+            # test_index_tricks.py#L513-L517
+            yield np.arange(16).reshape((4, 4))
+
+        for arr in inputs():
+            self.assertPreciseEqual(np.asarray(pyfunc(arr)), cfunc(arr))
+
+    def test_diag_indices_from_exception(self):
+        pyfunc = diag_indices_from
+        cfunc = njit(pyfunc)
+
+        self.disable_leak_check()
+
+        with self.assertRaises(TypingError) as raises:
+            cfunc("abc")
+        self.assertIn('The argument "arr" must be an array',
+                      str(raises.exception))
+
+        with self.assertRaises(TypingError) as raises:
+            cfunc("abc")
+        self.assertIn('The argument "arr" must be an array',
+                      str(raises.exception))
+
+        # Based on tests from https://github.com/numpy/numpy/blob/...
+        # 09ee7daa9113ecb503552913d76185d0ad3f0f31/numpy/lib/tests/...
+        # test_index_tricks.py#L519-L527
+        with self.assertRaises(ValueError) as raises:
+            cfunc(np.ones(7))
+        self.assertIn('Input array must be at least 2-d',
+                      str(raises.exception))
+
+        with self.assertRaises(ValueError) as raises:
+            cfunc(np.zeros((3, 3, 2, 3)))
+        self.assertIn('All dimensions of input must be of equal length',
+                      str(raises.exception))
+
 
 class TestNPMachineParameters(TestCase):
     # tests np.finfo, np.iinfo, np.MachAr