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)