Request: subsetting with numpy arrays/masks

[david-cortes-intel]

ref uxlfoundation/scikit-learn-intelex#2350

Currently, it's not possible to subset dpctl or dpnp arrays with a numpy array of integers:

import numpy as np
import dpnp
X = np.arange(16).reshape((4,4))
dpnp.array(X)[np.arange(2)]

---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
Cell In[9], line 4
      2 import dpnp
      3 X = np.arange(16).reshape((4,4))
----> 4 dpnp.array(X)[np.arange(2)]

File /localdisk2/mkl/dcortes/miniforge3/envs/icxconda/lib/python3.12/site-packages/dpnp/dpnp_array.py:355, in dpnp_array.__getitem__(self, key)
    352 """Return ``self[key]``."""
    353 key = _get_unwrapped_index_key(key)
--> 355 item = self._array_obj.__getitem__(key)
    356 return dpnp_array._create_from_usm_ndarray(item)

File dpctl/tensor/_usmarray.pyx:937, in dpctl.tensor._usmarray.usm_ndarray.__getitem__()

File dpctl/tensor/_slicing.pxi:300, in dpctl.tensor._usmarray._basic_slice_meta()

IndexError: Only integers, slices (`:`), ellipsis (`...`), dpctl.tensor.newaxis (`None`) and integer and boolean arrays are valid indices.

Allowing this type of indexing, even if not very efficiently, would be very helpful for scikit-learn-intelex.

This library scikit-learn-intelex is meant to be compatible with the library scikit-learn, by being able to use classes from scikit-learn and scikit-learn-intelex interchangeably, but with scikit-learn-intelex additionally being able to work with dpctl arrays and offering algorithms that run on GPU.

Internally, library scikit-learn (on which scikit-learn-intelex relies) has many functionalities which simply create arrays of indices to subset a larger array, which currently do not work with dpctl/dpnp inputs, as the integer indices are from numpy and the library is not sycl-aware.

Enabling subsetting of dpctl arrays with numpy integer and boolean arrays (as scikit-learn does internally) would immediately enable a lot of useful GPU features on scikit-learn-intelex, such as tuning parameters of machine learning models on GPU, calculating cross-validated metrics, among many others, which currently are not possible.

[ndgrigorian]

@david-cortes-intel
Is sending the array to the GPU via asarray insufficient for this purpose? i.e., is this a convenience feature? Because tensor.usm_ndarray are supported, and can be constructed from NumPy arrays via asarray or from_numpy

[david-cortes-intel]

@david-cortes-intel Is sending the array to the GPU via asarray insufficient for this purpose? i.e., is this a convenience feature? Because tensor.usm_ndarray are supported, and can be constructed from NumPy arrays via asarray or from_numpy

This is about compatibility with other tools in the ecosystem.

It might possible to send the array of indices to a device before subsetting it if you have code that is specifically designed for dpctl/dpnp.

But what we'd like to achieve here is being able to pass those dpnp/dpctl arrays to a library that is not specifically designed to handle those inputs and still have it work, even if this data movement makes the procedure less than optimal.

An example use-case is our library scikit-learn-intelex: it is meant to work with the library scikit-learn, which among many things, offers (a) the possibility of creating composable machine learning pipelines (with steps); and (b) the possibility of mixing scikit-learn-compatible components from different libraries in those pipelines, which typically handle only NumPy and Pandas inputs.

For example, a typical scikit-learn pipeline might involve the following steps:

1. Take two large arrays ('X' and 'y') and create multiple random subsets of it. This is simple python code that draws arrays of random numbers that are then used to subset other arrays.
2. Execute multiple runs of machine learning algorithms (with different parameters each) on those subsets. These are typically longer BLAS/LAPACK-based routines or optimized C procedures.
3. Calculate quality metrics from the machine learning algorithm (again C routines or vectorized numpy ops) on other subsets of the data.
4. Aggregate the results and output them to the user.

The scikit-learn-intelex library offers a function 'patch_sklearn' which replaces steps 2. and 3. above with calls to the oneDAL library, which has support for GPU algorithms and can work with dpctl/dpnp inputs, among others.

But while currently it's possible to execute those steps 2 and 3 independently with scikit-learn-intelex using GPU-allocated arrays as inputs, it's not possible to execute a whole pipeline on GPU arrays, because step 1 fails due to trying to subset with numpy arrays.

The only current workaround involves using NumPy arrays as the base input and then transferring the subsets to the device within each call to 2. and 3., which is much worse than transferring just the indices to device in step 1.

[ndgrigorian]

But what we'd like to achieve here is being able to pass those dpnp/dpctl arrays to a library that is not specifically designed to handle those inputs and still have it work, even if this data movement makes the procedure less than optimal.

Ok, I think I understand now, let me just clarify a bit:

you have some function foo(x, ...) where x is some array-like object, preferably an arbitrary one.

Somewhere in foo you have inds which is some random NumPy array of indices, and you want x[inds] to work because you want this function to also work for other, arbitrary array libraries. Interesting problem.

In general, this is what our array API compatibility is (theoretically) all about, but I do see your problem: neither dpctl nor dpnp offer the tools needed for the random input generation. So you're stuck with NumPy.

I'm not against the idea of NumPy arrays being automatically converted (when used as indices) by dpctl.tensor. dpnp may have also considered implementing this at some point.

@antonwolfy thoughts on this? Has dpnp ever considered allowing NumPy array indices? Does CuPy do this? __getitem__ could probably be overloaded if we don't implement it here.

[antonwolfy]

Currently dpnp fully leverages on dpctl behavior here. While CuPy forces array-like input indices to be converted to cupy.ndarray if possible.

In general, I see no issue with adding support for the same logic as CuPy has in dpnp at least, and to support transferring host data with indices to the device (inds will be allocated on the same sycl_queue and will inherit the same usm_type as input array x).

I guess, previously we don't have any use real case where it might be needed.

[antonwolfy]

One more request to support that raised in dpnp#2461.

[abagusetty]

Just wondering if there is any newer info on this effort

[ndgrigorian]

Just wondering if there is any newer info on this effort

Yes, work is starting on issue topic now that the 0.20.0/0.20.1 release is completed.

[ndgrigorian]

@abagusetty @david-cortes-intel
pinging to let you know that #2128 is now open handling this feature (still a draft—needs tests and probably a few changes)

[ndgrigorian]

The feature is now implemented, so this issue was closed