pydata / pydata/xarray

Error using vectorized indexing with array API compliant class

Open
#8,667 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

array API standard bug topic-arrays
Dominant language
Python
Stars
4.2k
Forks
1.4k
Avg merge
2d 15h
Merged PRs (30d)
14

Description

What happened?

Vectorized indexing can fail for array types that strictly follow the array API standard.

What did you expect to happen?

Vectorized indexing to all work.

Minimal Complete Verifiable Example
import numpy.array_api as nxp

da = xr.DataArray(
    nxp.reshape(nxp.arange(12), (3, 4)),
    dims=["x", "y"],
    coords={"x": [0, 1, 2], "y": ["a", "b", "c", "d"]},
)

da[[0, 2, 2], [1, 3]]  # works

ind_x = xr.DataArray([0, 1], dims=["x"])
ind_y = xr.DataArray([0, 1], dims=["y"])

da[ind_x, ind_y]  # works

da[[0, 1], ind_x]  # doesn't work

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[157], line 1
----> 1 da[[0, 1], ind_x]

File ~/Documents/Work/Code/xarray/xarray/core/dataarray.py:859, in DataArray.__getitem__(self, key)
    856     return self._getitem_coord(key)
    857 else:
    858     # xarray-style array indexing
--> 859     return self.isel(indexers=self._item_key_to_dict(key))

File ~/Documents/Work/Code/xarray/xarray/core/dataarray.py:1472, in DataArray.isel(self, indexers, drop, missing_dims, **indexers_kwargs)
   1469 indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "isel")
   1471 if any(is_fancy_indexer(idx) for idx in indexers.values()):
-> 1472     ds = self._to_temp_dataset()._isel_fancy(
   1473         indexers, drop=drop, missing_dims=missing_dims
   1474     )
   1475     return self._from_temp_dataset(ds)
   1477 # Much faster algorithm for when all indexers are ints, slices, one-dimensional
   1478 # lists, or zero or one-dimensional np.ndarray's

File ~/Documents/Work/Code/xarray/xarray/core/dataset.py:3001, in Dataset._isel_fancy(self, indexers, drop, missing_dims)
   2997 var_indexers = {
   2998     k: v for k, v in valid_indexers.items() if k in var.dims
   2999 }
   3000 if var_indexers:
-> 3001     new_var = var.isel(indexers=var_indexers)
   3002     # drop scalar coordinates
   3003     # https://github.com/pydata/xarray/issues/6554
   3004     if name in self.coords and drop and new_var.ndim == 0:

File ~/Documents/Work/Code/xarray/xarray/core/variable.py:1130, in Variable.isel(self, indexers, missing_dims, **indexers_kwargs)
   1127 indexers = drop_dims_from_indexers(indexers, self.dims, missing_dims)
   1129 key = tuple(indexers.get(dim, slice(None)) for dim in self.dims)
-> 1130 return self[key]

File ~/Documents/Work/Code/xarray/xarray/core/variable.py:812, in Variable.__getitem__(self, key)
    799 """Return a new Variable object whose contents are consistent with
    800 getting the provided key from the underlying data.
    801 
   (...)
    809 array `x.values` directly.
    810 """
    811 dims, indexer, new_order = self._broadcast_indexes(key)
--> 812 data = as_indexable(self._data)[indexer]
    813 if new_order:
    814     data = np.moveaxis(data, range(len(new_order)), new_order)

File ~/Documents/Work/Code/xarray/xarray/core/indexing.py:1390, in ArrayApiIndexingAdapter.__getitem__(self, key)
   1388 else:
   1389     if isinstance(key, VectorizedIndexer):
-> 1390         raise TypeError("Vectorized indexing is not supported")
   1391     else:
   1392         raise TypeError(f"Unrecognized indexer: {key}")

TypeError: Vectorized indexing is not supported
MVCE confirmation
  • Minimal example — the example is as focused as reasonably possible to demonstrate the underlying issue in xarray.
  • Complete example — the example is self-contained, including all data and the text of any traceback.
  • Verifiable example — the example copy & pastes into an IPython prompt or Binder notebook, returning the result.
  • New issue — a search of GitHub Issues suggests this is not a duplicate.
  • Recent environment — the issue occurs with the latest version of xarray and its dependencies.
Relevant log output

No response

Anything else we need to know?

I don't really understand why the first two examples work but the last one doesn't...

Environment

main branch of xarray, numpy 1.26.0

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with xarray/core/indexing.py at ArrayApiIndexingAdapter.getitem and reproduce the failure using the provided minimal example with numpy.array_api. Trace how mixed list and DataArray indexers reach VectorizedIndexer handling; done means the failing expression works without the TypeError while the two working examples continue to work.

Written by the indexing model from the issue text.

Assessment

Tech stack
numpy, python
Domain
data
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.