Skip to content
Merged
Show file tree
Hide file tree
Changes from 8 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 0 additions & 35 deletions dpnp/dpnp_algo/dpnp_algo_indexing.pxi
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,6 @@ __all__ += [
"dpnp_diag_indices",
"dpnp_diagonal",
"dpnp_fill_diagonal",
"dpnp_indices",
"dpnp_putmask",
"dpnp_select",
"dpnp_tril_indices",
Expand Down Expand Up @@ -237,40 +236,6 @@ cpdef dpnp_fill_diagonal(dpnp_descriptor x1, val):
c_dpctl.DPCTLEvent_Delete(event_ref)


cpdef object dpnp_indices(dimensions):
len_dimensions = len(dimensions)
res_shape = []
res_shape.append(len_dimensions)
for i in range(len_dimensions):
res_shape.append(dimensions[i])

result = []
if len_dimensions == 1:
res = []
for i in range(dimensions[0]):
res.append(i)
result.append(res)
else:
res1 = []
for i in range(dimensions[0]):
res = []
for j in range(dimensions[1]):
res.append(i)
res1.append(res)
result.append(res1)

res2 = []
for i in range(dimensions[0]):
res = []
for j in range(dimensions[1]):
res.append(j)
res2.append(res)
result.append(res2)

dpnp_result = dpnp.array(result)
return dpnp_result


cpdef dpnp_putmask(utils.dpnp_descriptor arr, utils.dpnp_descriptor mask, utils.dpnp_descriptor values):
cdef int values_size = values.size

Expand Down
103 changes: 86 additions & 17 deletions dpnp/dpnp_iface_indexing.py
Original file line number Diff line number Diff line change
Expand Up @@ -355,31 +355,100 @@ def fill_diagonal(x1, val, wrap=False):
return call_origin(numpy.fill_diagonal, x1, val, wrap, dpnp_inplace=True)


def indices(dimensions, dtype=int, sparse=False):
def indices(
dimensions,
dtype=int,
sparse=False,
device=None,
usm_type="device",
sycl_queue=None,
):
"""
Return an array representing the indices of a grid.

Compute an array where the subarrays contain index values 0, 1, …
varying only along the corresponding axis.

For full documentation refer to :obj:`numpy.indices`.

Limitations
-----------
Parameters `dtype` and `sparse` are supported only with default values.
Parameter `dimensions` is supported with len <=2.
Parameters
----------
dimensions : sequence of ints
The shape of the grid.
dtype : dtype, optional
Data type of the result.
sparse : boolean, optional
Return a sparse representation of the grid instead of a dense representation.
Default is ``False``.

"""
Returns
-------
out : one dpnp.ndarray or tuple of dpnp.ndarray
If sparse is ``False``:
Returns one array of grid indices, grid.shape = (len(dimensions),) + tuple(dimensions).

if not isinstance(dimensions, (tuple, list)):
pass
elif len(dimensions) > 2 or len(dimensions) == 0:
pass
elif dtype != int:
pass
elif sparse:
pass
else:
return dpnp_indices(dimensions)
If sparse is ``True``:
Returns a tuple of arrays, with grid[i].shape = (1, ..., 1, dimensions[i], 1, ..., 1)
with dimensions[i] in the ith place.

return call_origin(numpy.indices, dimensions, dtype, sparse)
Examples
--------
>>> import dpnp as np
>>> grid = np.indices((2, 3))
>>> grid.shape
(2, 2, 3)
>>> grid[0]
array([[0, 0, 0],
[1, 1, 1]])
>>> grid[1]
array([[0, 1, 2],
[0, 1, 2]])

>>> x = np.arange(20).reshape(5, 4)
>>> row, col = np.indices((2, 3))
>>> x[row, col]
array([[0, 1, 2],
[4, 5, 6]])

>>> i, j = np.indices((2, 3), sparse=True)
>>> i.shape
(2, 1)
>>> j.shape
(1, 3)
>>> i
array([[0],
[1]])
>>> j
array([[0, 1, 2]])

"""

dimensions = tuple(dimensions)
N = len(dimensions)
shape = (1,) * N
if sparse:
res = ()
else:
res = dpnp.empty(
(N,) + dimensions,
dtype=dtype,
device=device,
usm_type=usm_type,
sycl_queue=sycl_queue,
)
for i, dim in enumerate(dimensions):
idx = dpnp.arange(
dim,
dtype=dtype,
device=device,
usm_type=usm_type,
sycl_queue=sycl_queue,
).reshape(shape[:i] + (dim,) + shape[i + 1 :])
if sparse:
res = res + (idx,)
else:
res[i] = idx
return res


def nonzero(x, /):
Expand Down
11 changes: 7 additions & 4 deletions tests/test_indexing.py
Original file line number Diff line number Diff line change
Expand Up @@ -363,10 +363,13 @@ def test_fill_diagonal(array, val):
"[3, 2]",
],
)
def test_indices(dimension):
expected = numpy.indices(dimension)
result = dpnp.indices(dimension)
assert_array_equal(expected, result)
@pytest.mark.parametrize("dtype", get_all_dtypes(no_bool=True))
@pytest.mark.parametrize("sparse", [True, False], ids=["True", "False"])
def test_indices(dimension, dtype, sparse):
expected = numpy.indices(dimension, dtype=dtype, sparse=sparse)
result = dpnp.indices(dimension, dtype=dtype, sparse=sparse)
for Xnp, X in zip(expected, result):
assert_array_equal(Xnp, X)


@pytest.mark.parametrize(
Expand Down
12 changes: 12 additions & 0 deletions tests/test_sycl_queue.py
Original file line number Diff line number Diff line change
Expand Up @@ -1418,6 +1418,18 @@ def test_take(func, device):
assert_sycl_queue_equal(result_queue, expected_queue)


@pytest.mark.parametrize(
"device",
valid_devices,
ids=[device.filter_string for device in valid_devices],
)
@pytest.mark.parametrize("sparse", [True, False], ids=["True", "False"])
def test_indices(device, sparse):
dpnp_array = dpnp.indices((2, 3), sparse=sparse, device=device)
for i in dpnp_array:
assert i.sycl_device == device


@pytest.mark.parametrize(
"device",
valid_devices,
Expand Down
8 changes: 8 additions & 0 deletions tests/test_usm_type.py
Original file line number Diff line number Diff line change
Expand Up @@ -521,6 +521,14 @@ def test_take(func, usm_type_x, usm_type_ind):
assert z.usm_type == du.get_coerced_usm_type([usm_type_x, usm_type_ind])


@pytest.mark.parametrize("usm_type", list_of_usm_types, ids=list_of_usm_types)
@pytest.mark.parametrize("sparse", [True, False], ids=["True", "False"])
def test_indices(usm_type, sparse):
x = dp.indices((2, 3), sparse=sparse, usm_type=usm_type)
for i in x:
assert i.usm_type == usm_type


@pytest.mark.parametrize(
"usm_type_matrix", list_of_usm_types, ids=list_of_usm_types
)
Expand Down
1 change: 0 additions & 1 deletion tests/third_party/cupy/indexing_tests/test_generate.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@
from tests.third_party.cupy import testing


@pytest.mark.usefixtures("allow_fall_back_on_numpy")
@testing.gpu
class TestIndices(unittest.TestCase):
@testing.for_all_dtypes()
Expand Down