Slicing with lists in multiple axes

[jhamman]

From the From the dask docs:

Dask.array supports most of the NumPy slicing syntax.
...
It does not currently support the following:

Slicing one dask.array with another x[x > 0]
Slicing with lists in multiple axes x[[1, 2, 3], [3, 2, 1]]

Both of these are straightforward to add though. If you have a use case then raise an issue.

Here's that issue.

My use case is for point-wise indexing in xray: pydata/xarray#475

A simple use case using dask arrays:

x = da.ones((10, 100), chunks=(10, 10))
points = x[[1, 2, 3], [3, 2, 1]]

currently raises this error:

NotImplementedError: Don't yet support nd fancy indexing

cc @shoyer

[mrocklin]

Am I correct in assuming that in this case you'd like the shape of the output to be (3,) and not (3, 3)?

[mrocklin]

And if so is it important to you that this be the default array slicing syntax or are you ok with some other custom method?

[jhamman]

Correct. I'm looking for the default numpy slicing behavior, how we get there is up for discussion. My preference would be to use the numpy slicing syntax, but could use a take method of some kind if that is substantially easier.

[mrocklin]

They're both the same level of difficulty to accomplish. I mostly want to avoid locking dask into one behavior or the other for as long as possible.

[shoyer]

NumPy has been discussing adding vindex and oindex attributes for explicit vectorized vs. outer/orthogonal indexing (see @seberg's NEP and PR).

I suggest that dask should follow NumPy's lead here, and consider implementing both attributes -- even if it will only support a limited subset of vectorized indexing. Outer indexing is easier to reason about and optimize (dask already has all the necessary functionality), and it will be useful to have explicit syntax both to support better fusing of getitem calls and to make things simpler for downstream libraries like xray (which uses outer indexing in __getitem__).

Dask certainly should not try to exactly replicate NumPy's current indexing behavior, which sometimes but not always reorders axes with array indices, e.g.,

In [14]: x = np.zeros((5, 6, 7, 8))

In [15]: x[:, [0, 1], [0, 1]].shape
Out[15]: (5, 2, 8)

In [16]: x[:, [0, 1], :, [0, 1]].shape
Out[16]: (2, 5, 7)

[mrocklin]

OK, I've implemented this (I think) in #439 . It could use a better name. Happy to use vindex if that's best. Feedback welcome.