Skip to content

Commit f703070

Browse files
committed
Leverage dpnp.ndindex on numpy class
1 parent c6c478a commit f703070

File tree

1 file changed

+52
-0
lines changed

1 file changed

+52
-0
lines changed

dpnp/dpnp_iface_indexing.py

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,7 @@
6464
"indices",
6565
"ix_",
6666
"mask_indices",
67+
"ndindex",
6768
"nonzero",
6869
"place",
6970
"put",
@@ -1057,6 +1058,57 @@ def mask_indices(
10571058
return nonzero(a != 0)
10581059

10591060

1061+
# pylint: disable=invalid-name
1062+
ndindex = numpy.ndindex
1063+
ndindex.__doc__ = """
1064+
An N-dimensional iterator object to index arrays.
1065+
1066+
Given the shape of an array, an `ndindex` instance iterates over the
1067+
N-dimensional index of the array. At each iteration a tuple of indices is
1068+
returned, the last dimension is iterated over first.
1069+
1070+
For full documentation refer to :obj:`numpy.ndindex`.
1071+
1072+
Parameters
1073+
----------
1074+
shape : ints, or a single tuple of ints
1075+
The size of each dimension of the array can be passed as individual
1076+
parameters or as the elements of a tuple.
1077+
1078+
See Also
1079+
--------
1080+
:obj:`dpnp.ndenumerate` : Multidimensional index iterator.
1081+
:obj:`dpnp.flatiter` : Flat iterator object to iterate over arrays.
1082+
1083+
Examples
1084+
--------
1085+
>>> import dpnp as np
1086+
1087+
Dimensions as individual arguments
1088+
1089+
>>> for index in np.ndindex(3, 2, 1):
1090+
... print(index)
1091+
(0, 0, 0)
1092+
(0, 1, 0)
1093+
(1, 0, 0)
1094+
(1, 1, 0)
1095+
(2, 0, 0)
1096+
(2, 1, 0)
1097+
1098+
Same dimensions - but in a tuple ``(3, 2, 1)``
1099+
1100+
>>> for index in np.ndindex((3, 2, 1)):
1101+
... print(index)
1102+
(0, 0, 0)
1103+
(0, 1, 0)
1104+
(1, 0, 0)
1105+
(1, 1, 0)
1106+
(2, 0, 0)
1107+
(2, 1, 0)
1108+
1109+
"""
1110+
1111+
10601112
def nonzero(a):
10611113
"""
10621114
Return the indices of the elements that are non-zero.

0 commit comments

Comments
 (0)