Indexing `dpnp.ndarray` indexing with numpy lists-of-lists

[abagusetty]

It looks like I have hit an other corner case with indexing dpnp.ndarray with numpy indices. There has already been recent support for this and was wondering if someone look into this.

The following reproducer was extracted from an application with mimicking cupy's usage

import numpy as np
import dpnp as dpnp
import dpctl

print("dpnp:", dpnp.__version__)
print("dpctl:", dpctl.__version__)

# Simulate frr with shape (3, ngrids). Content is irrelevant; shape matters.
ngrids = 10240
frr = dpnp.arange(3 * ngrids, dtype=dpnp.float64).reshape(3, ngrids)
print("frr:", type(frr), frr.shape)

# This is done via a list-of-lists assigned as the index for one axis:
# slices[axis] = [[0, 1], [1, 2]]
def stack_frr_like(frr, axis=0):
    slices = [slice(None)] * frr.ndim
    # --- Case A: list-of-lists (exactly like PySCF) ---
    slices[axis] = [[0, 1], [1, 2]]
    print("\nCase A: list-of-lists index:", slices[axis], "(type:", type(slices[axis]), ")")
    try:
        out = frr[tuple(slices)]
        print("Case A: OK, out.shape =", out.shape)
    except Exception as e:
        print("Case A: FAILED with:", repr(e))

    # --- Case B: NumPy integer array index ---
    slices = [slice(None)] * frr.ndim
    slices[axis] = np.array([[0, 1], [1, 2]], dtype=np.intp)
    print("\nCase B: NumPy array index dtype:", slices[axis].dtype, "(type:", type(slices[axis]), ")")
    try:
        out = frr[tuple(slices)]
        print("Case B: OK, out.shape =", out.shape)
    except Exception as e:
        print("Case B: FAILED with:", repr(e))

    # --- Case C: dpnp (device) integer array index (WORKAROUND) ---
    slices = [slice(None)] * frr.ndim
    slices[axis] = dpnp.asarray(np.array([[0, 1], [1, 2]], dtype=np.intp))
    print("\nCase C: dpnp array index dtype:", slices[axis].dtype, "(type:", type(slices[axis]), ")")
    try:
        out = frr[tuple(slices)]
        print("Case C: OK, out.shape =", out.shape)
    except Exception as e:
        print("Case C: FAILED with:", repr(e))

stack_frr_like(frr, axis=0)

[antonwolfy]

@ndgrigorian, do we have any plans to support a list as valid indices in dpctl?