From f0b183776991b795bc5b77b66a746f925acfbd4d Mon Sep 17 00:00:00 2001 From: pfackeldey Date: Fri, 13 Dec 2024 16:57:33 -0500 Subject: [PATCH 1/2] allow int as placeholder array index type --- src/awkward/_nplikes/placeholder.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/awkward/_nplikes/placeholder.py b/src/awkward/_nplikes/placeholder.py index 20a2af6ee7..94bd3dd98f 100644 --- a/src/awkward/_nplikes/placeholder.py +++ b/src/awkward/_nplikes/placeholder.py @@ -93,6 +93,8 @@ def __getitem__(self, index): new_length = (stop - start) // step return type(self)(self._nplike, (new_length,), self._dtype) + elif isinstance(index, int): + return type(self)(self._nplike, (1,), self._dtype) else: raise TypeError( f"{type(self).__name__} supports only trivial slices, not {type(index).__name__}" From 647ab310e71057629b1b4bb1c058e01c67c2d038 Mon Sep 17 00:00:00 2001 From: pfackeldey Date: Fri, 13 Dec 2024 17:59:42 -0500 Subject: [PATCH 2/2] also add indexing with array of integers --- src/awkward/_nplikes/placeholder.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/awkward/_nplikes/placeholder.py b/src/awkward/_nplikes/placeholder.py index 94bd3dd98f..cff0698e82 100644 --- a/src/awkward/_nplikes/placeholder.py +++ b/src/awkward/_nplikes/placeholder.py @@ -8,6 +8,7 @@ from awkward._nplikes.array_like import ArrayLike from awkward._nplikes.numpy_like import NumpyLike, NumpyMetadata from awkward._nplikes.shape import ShapeItem, unknown_length +from awkward._regularize import is_array_like from awkward._typing import TYPE_CHECKING, Any, DType, Self np = NumpyMetadata.instance() @@ -95,6 +96,8 @@ def __getitem__(self, index): return type(self)(self._nplike, (new_length,), self._dtype) elif isinstance(index, int): return type(self)(self._nplike, (1,), self._dtype) + elif is_array_like(index) and np.issubdtype(index.dtype, np.integer): + return type(self)(self._nplike, index.shape, self._dtype) else: raise TypeError( f"{type(self).__name__} supports only trivial slices, not {type(index).__name__}"