Skip to content

Support buffer protocol objects as advanced index keys #2889

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
vlad-perevezentsev wants to merge 3 commits into
base: fix_inplace_indexind_4d
Choose a base branch
from
+53 −4
Open

Support buffer protocol objects as advanced index keys #2889

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
4523353
Support buffer protocol objects as index keys
vlad-perevezentsev May 4, 2026
fc5b54f
Add tests for buffer protocol advanced indexing
vlad-perevezentsev May 4, 2026
c883880
Update changelog
vlad-perevezentsev May 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions .github/workflows/conda-package.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -91,13 +91,13 @@ jobs:
- name: Build conda package
id: build_conda_pkg
continue-on-error: true
run: conda build --no-test --python ${{ matrix.python }} --numpy 2.0 ${{ env.channels-list }} conda-recipe
run: conda-build --no-test --python ${{ matrix.python }} --numpy 2.0 ${{ env.channels-list }} conda-recipe
env:
MAX_BUILD_CMPL_MKL_VERSION: '2026.0a0'

- name: ReBuild conda package
if: steps.build_conda_pkg.outcome == 'failure'
run: conda build --no-test --python ${{ matrix.python }} --numpy 2.0 ${{ env.channels-list }} conda-recipe
run: conda-build --no-test --python ${{ matrix.python }} --numpy 2.0 ${{ env.channels-list }} conda-recipe
env:
MAX_BUILD_CMPL_MKL_VERSION: '2026.0a0'

Expand Down
2 changes: 2 additions & 0 deletions CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0

### Added

* Added support for buffer protocol objects as advanced index keys in `dpnp.ndarray` [#2889](https://github.com/IntelPython/dpnp/pull/2889)

### Changed

### Deprecated
Expand Down
15 changes: 13 additions & 2 deletions dpnp/dpnp_array.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,8 +52,9 @@ def _unwrap_index_element(x):
"""
Unwrap a single index element for the tensor indexing layer.

Converts dpnp arrays to usm_ndarray and array-like objects (range, list)
to numpy arrays with intp dtype for NumPy-compatible advanced indexing.
Converts dpnp arrays to usm_ndarray and array-like objects (range, list,
buffer protocol objects) to numpy arrays for NumPy-compatible advanced
indexing.

"""

Expand All @@ -71,6 +72,16 @@ def _unwrap_index_element(x):
if arr.size == 0:
arr = arr.astype(numpy.intp)
return arr
if isinstance(x, numpy.ndarray):
return x
# convert buffer protocol objects (array.array, memoryview, etc.)
Copy link
Copy Markdown
Contributor

@antonwolfy antonwolfy May 5, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can all that use cases covered converting x to numpy.ndarray?
Do we need so complicated if-logic?

Copy link
Copy Markdown
Collaborator

@ndgrigorian ndgrigorian May 5, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

here is CuPy implementation of a similar idea for a reference:
https://github.com/cupy/cupy/blob/main/cupy/_core/_routines_indexing.pyx#L247

can take a similar approach too in some cases (best to leave converting to usm_ndarray to the tensor layer, though, for queue propagation reasons)

Copy link
Copy Markdown
Contributor

@antonwolfy antonwolfy May 5, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think we need to fully follow the same CuPy implementation in _prepare_slice_list, because it seems CuPy mimics faulty NumPy behavior sometimes.

Not sure if I got the comment fully correctly, but in general I meant something like:

    if x is None or x is Ellipsis or isinstance(x, (dpt.usm_ndarray, slice, numpy.ndarray)):
        return x
    if isinstance(x, dpnp_array):
        return x.get_array()

    try:
        x = numpy.asarray(x)
    except ... :
        ...
        raise ...
    ...
    return x

Copy link
Copy Markdown
Collaborator

@ndgrigorian ndgrigorian May 5, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yep, agreed mostly, though problem is, we don't really want to send integers or boolean Python scalars to numpy arrays, which is why we may need something like:

        elif numpy.isscalar(s):
            if not isinstance(s, (bool, numpy.bool_)):
                # keep scalar int
                continue

which is in cupy implementation. We could handle range, list, tuple through numpy though

try:
Copy link
Copy Markdown
Collaborator

@ndgrigorian ndgrigorian May 4, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

seems __buffer__ has been added with 3.12: https://docs.python.org/3/reference/datamodel.html#object.__buffer__
https://docs.python.org/3/library/collections.abc.html#collections.abc.Buffer
https://peps.python.org/pep-0688/

we should probably start to support it as well. Maybe add a TODO to check for __buffer__/use collections.abc.Buffer and then call memoryview if it fails

Copy link
Copy Markdown
Collaborator

@ndgrigorian ndgrigorian May 5, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

actually I misunderstood the PEP: this feature is already implemented, we can add a check for Python version and check against __buffer__ or try memoryview only if it's <3.12

mv = memoryview(x)
except TypeError:
return x
# 0-d buffers are handled by the tensor layer
if mv.ndim > 0:
return numpy.asarray(x)
return x


Expand Down
36 changes: 36 additions & 0 deletions dpnp/tests/test_indexing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import array
import functools

import dpctl
Expand Down Expand Up @@ -406,6 +407,41 @@ def test_array_like_single_index(self, idx):
dp_a = dpnp.arange(24).reshape(2, 3, 4)
assert_array_equal(dp_a[idx], np_a[idx])

def test_buffer_protocol_getitem(self):
inds = array.array("l")
inds.frombytes(numpy.arange(3).tobytes())
np_a = numpy.arange(12).reshape(3, 4)
dp_a = dpnp.arange(12).reshape(3, 4)
assert_array_equal(dp_a[inds], np_a[inds])

def test_buffer_protocol_paired_index(self):
inds = array.array("l")
inds.frombytes(numpy.arange(3).tobytes())
np_a = numpy.arange(12).reshape(3, 4)
dp_a = dpnp.arange(12).reshape(3, 4)
assert_array_equal(dp_a[inds, inds], np_a[inds, inds])

def test_buffer_protocol_setitem(self):
inds = array.array("l")
inds.frombytes(numpy.arange(3).tobytes())
np_a = numpy.arange(12).reshape(3, 4)
dp_a = dpnp.arange(12).reshape(3, 4)
np_a[inds, inds] = 0
dp_a[inds, inds] = 0
assert_array_equal(dp_a, np_a)

def test_memoryview_getitem(self):
inds = memoryview(array.array("l", [0, 1, 2]))
np_a = numpy.arange(12).reshape(3, 4)
dp_a = dpnp.arange(12).reshape(3, 4)
assert_array_equal(dp_a[inds], np_a[inds])

def test_bytearray_getitem(self):
inds = bytearray(b"\x00\x01\x02")
np_a = numpy.arange(10)
dp_a = dpnp.arange(10)
assert_array_equal(dp_a[inds], np_a[inds])


class TestIx:
@pytest.mark.parametrize(
Expand Down
Loading