Merge d77a69f into d06277f

Author: antonwolfy

dpnp/dpnp_iface_indexing.py

@@ -67,6 +67,7 @@
     "indices",
     "ix_",
     "mask_indices",
+    "ndindex",
     "nonzero",
     "place",
     "put",
@@ -1059,6 +1060,57 @@ def mask_indices(
     return nonzero(a != 0)
 
 
+# pylint: disable=invalid-name
+ndindex = numpy.ndindex
+ndindex.__doc__ = """
+An N-dimensional iterator object to index arrays.
+
+Given the shape of an array, an `ndindex` instance iterates over the
+N-dimensional index of the array. At each iteration a tuple of indices is
+returned, the last dimension is iterated over first.
+
+For full documentation refer to :obj:`numpy.ndindex`.
+
+Parameters
+----------
+shape : ints, or a single tuple of ints
+    The size of each dimension of the array can be passed as individual
+    parameters or as the elements of a tuple.
+
+See Also
+--------
+:obj:`dpnp.ndenumerate` : Multidimensional index iterator.
+:obj:`dpnp.flatiter` : Flat iterator object to iterate over arrays.
+
+Examples
+--------
+>>> import dpnp as np
+
+Dimensions as individual arguments
+
+>>> for index in np.ndindex(3, 2, 1):
+...     print(index)
+(0, 0, 0)
+(0, 1, 0)
+(1, 0, 0)
+(1, 1, 0)
+(2, 0, 0)
+(2, 1, 0)
+
+Same dimensions - but in a tuple ``(3, 2, 1)``
+
+>>> for index in np.ndindex((3, 2, 1)):
+...     print(index)
+(0, 0, 0)
+(0, 1, 0)
+(1, 0, 0)
+(1, 1, 0)
+(2, 0, 0)
+(2, 1, 0)
+
+"""
+
+
 def nonzero(a):
     """
     Return the indices of the elements that are non-zero.

dpnp/dpnp_utils/dpnp_utils_pad.py

@@ -642,8 +642,7 @@ def dpnp_pad(array, pad_width, mode="constant", **kwargs):
 
             # compute indices for the iteration axes, and append a trailing
             # ellipsis to prevent 0d arrays decaying to scalars
-            # TODO: replace with dpnp.ndindex when implemented
-            inds = numpy.ndindex(view.shape[:-1])
+            inds = dpnp.ndindex(view.shape[:-1])
             inds = (ind + (Ellipsis,) for ind in inds)
             for ind in inds:
                 function(view[ind], pad_width[axis], axis, kwargs)

tests/test_indexing.py

@@ -335,6 +335,18 @@ def test_ix_error(self, xp, shape):
         assert_raises(ValueError, xp.ix_, xp.ones(shape))
 
 
+class TestNdindex:
+    @pytest.mark.parametrize(
+        "shape", [[1, 2, 3], [(1, 2, 3)], [(3,)], [3], [], [()], [0]]
+    )
+    def test_basic(self, shape):
+        result = dpnp.ndindex(*shape)
+        expected = numpy.ndindex(*shape)
+
+        for x, y in zip(result, expected):
+            assert x == y
+
+
 class TestNonzero:
     @pytest.mark.parametrize("list_val", [[], [0], [1]])
     def test_trivial(self, list_val):