[dpnp.array] Indexing issue

[abagusetty]

I was trying to port a CUPY snippet for ALCF Aurora and was experiencing issue with DPNP array indexing with np.array, etc. Does someone have any suggestions. I could flatten the indexing-array but prefer not to since the usage is at several places. Appreciate any pointers.

Reproducer:

import dpnp as dp
import numpy as np

# Create a 2D dpnp array
sorted_mat = dp.arange(9).reshape(3, 3)

# Simulate a fancy index like your reshaped AO index
idx = np.array([2, 0, 1])
fancy_index = [idx.reshape(-1, 1), np.arange(3).reshape(1, -1)]

# Try to apply NumPy-style fancy indexing
try:
    mat = dp.empty_like(sorted_mat)
    mat[fancy_index[0], fancy_index[1]] = sorted_mat
except Exception as e:
    print("Expected failure:")
    print(type(e).__name__ + ":", e)

[antonwolfy]

@abagusetty, thank you for the reporting issue.
Is that intended to pass numpy.ndarray as a fancy index?
Currently only dpnp.ndarray, integers, slices (:), ellipsis (...), dpnp.newaxis (None) are supported there.

[abagusetty]

@antonwolfy Appreciate the prompt response. Yes, It looks like this line in the original snippet indexes as a numpy.ndarray for the cupy array. https://github.com/pyscf/gpu4pyscf/blob/master/gpu4pyscf/gto/int3c1e.py#L193

mat[tuple(fancy_index)] = sorted_mat

We were just thinking, if we had to change the original source from indexing numpy.ndarray to dpnp.ndarray. Currently we are monkey patching some equivalents of cupy in __init__.py such that every occurrence of cupy maps to dpnp equivalents. And many occurrences have worked well except this one.

[abagusetty]

Sorry I should have been explicit, the fancy_index in the context is a list of numpy.ndarray

[antonwolfy]

Here is open issue dpctl#2053 to support that use case.
Once implemented in dpctl, the above example will work with dpnp.

[abagusetty]

@antonwolfy Thanks so much, will keep an eye and try out the solution. Appreciate the super quick help.