From b18a02914d55a1783b379b3fa35044c6c36ae7ff Mon Sep 17 00:00:00 2001 From: Ianna Osborne Date: Mon, 16 Sep 2024 16:20:18 +0200 Subject: [PATCH 1/2] debug: slicing on gpu --- src/awkward/_slicing.py | 28 +- src/awkward/contents/content.py | 29 + src/awkward/contents/listarray.py | 23 + src/awkward/contents/listoffsetarray.py | 23 + src/awkward/contents/numpyarray.py | 21 + tests-cuda/test_3140_cuda_slicing.py | 1349 ++++++++++++----------- 6 files changed, 804 insertions(+), 669 deletions(-) diff --git a/src/awkward/_slicing.py b/src/awkward/_slicing.py index 9f6fbe3741..f8193951be 100644 --- a/src/awkward/_slicing.py +++ b/src/awkward/_slicing.py @@ -21,7 +21,19 @@ SliceItem: TypeAlias = "int | slice | str | None | Ellipsis | ArrayLike | Content" +import functools +def trace_function_calls(func): + @functools.wraps(func) + def wrapper(*args, **kwargs): + print(f"_slicing.py: Calling function: {func.__name__}") + result = func(*args, **kwargs) + print(f"_slicing.py: Function {func.__name__} returned {result}") + return result + return wrapper + + +@trace_function_calls def normalize_slice(slice_: slice, *, nplike: NumpyLike) -> slice: """ Args: @@ -59,6 +71,7 @@ def __repr__(self): S = TypeVar("S", bound=Sequence) +@trace_function_calls def head_tail(sequence: S[T]) -> tuple[T | type(NO_HEAD), S[T]]: if len(sequence) == 0: return NO_HEAD, () @@ -66,6 +79,7 @@ def head_tail(sequence: S[T]) -> tuple[T | type(NO_HEAD), S[T]]: return sequence[0], sequence[1:] +@trace_function_calls def prepare_advanced_indexing(items, backend: Backend): """Broadcast index objects to satisfy NumPy indexing rules @@ -177,7 +191,7 @@ def prepare_advanced_indexing(items, backend: Backend): ) return tuple(prepared) - +@trace_function_calls def normalize_integer_like(x) -> int | ArrayLike: if is_array_like(x): if np.issubdtype(x.dtype, np.integer) and x.ndim == 0: @@ -187,7 +201,7 @@ def normalize_integer_like(x) -> int | ArrayLike: else: return int(x) - +@trace_function_calls def normalise_item(item, backend: Backend) -> SliceItem: """ Args: @@ -300,12 +314,12 @@ def normalise_item(item, backend: Backend) -> SliceItem: + repr(item).replace("\n", "\n ") ) - +@trace_function_calls def normalise_items(where: Sequence, backend: Backend) -> list: # First prepare items for broadcasting into like-types return [normalise_item(x, backend=backend) for x in where] - +@trace_function_calls def _normalise_item_RegularArray_to_ListOffsetArray64(item: Content) -> Content: if isinstance(item, ak.contents.RegularArray): next = item.to_ListOffsetArray64() @@ -321,7 +335,7 @@ def _normalise_item_RegularArray_to_ListOffsetArray64(item: Content) -> Content: else: raise AssertionError(type(item)) - +@trace_function_calls def _normalise_item_nested(item: Content) -> Content: if isinstance(item, ak.contents.EmptyArray): # policy: unknown -> int @@ -460,7 +474,7 @@ def _normalise_item_nested(item: Content) -> Content: + repr(item).replace("\n", "\n ") ) - +@trace_function_calls def _normalise_item_bool_to_int(item: Content, backend: Backend) -> Content: """ Args: @@ -650,7 +664,7 @@ def _normalise_item_bool_to_int(item: Content, backend: Backend) -> Content: else: raise AssertionError(type(item)) - +@trace_function_calls def getitem_next_array_wrap( outcontent: Content, shape: tuple[int], outer_length: int = 0 ) -> Content: diff --git a/src/awkward/contents/content.py b/src/awkward/contents/content.py index d0169ee2eb..5d71f5064e 100644 --- a/src/awkward/contents/content.py +++ b/src/awkward/contents/content.py @@ -72,7 +72,16 @@ JSONValueType: TypeAlias = """ float | int | str | list[JSONValueType] | dict[str, JSONValueType] """ +import functools +def trace_function_calls(func): + @functools.wraps(func) + def wrapper(*args, **kwargs): + print(f"content.py: Calling function: {func.__name__}") + result = func(*args, **kwargs) + print(f"content.py: Function {func.__name__} returned {result}") + return result + return wrapper class ImplementsApplyAction(Protocol): def __call__( @@ -298,6 +307,7 @@ def __iter__(self): for i in range(len(self)): yield self._getitem_at(i) + @trace_function_calls def _getitem_next_field( self, head: SliceItem | tuple, @@ -307,6 +317,7 @@ def _getitem_next_field( nexthead, nexttail = ak._slicing.head_tail(tail) return self._getitem_field(head)._getitem_next(nexthead, nexttail, advanced) + @trace_function_calls def _getitem_next_fields( self, head: SliceItem, tail: tuple[SliceItem, ...], advanced: Index | None ) -> Content: @@ -321,6 +332,8 @@ def _getitem_next_fields( nexthead, nexttail, advanced ) + + @trace_function_calls def _getitem_next_newaxis( self, tail: tuple[SliceItem, ...], advanced: Index | None ) -> RegularArray: @@ -329,6 +342,7 @@ def _getitem_next_newaxis( self._getitem_next(nexthead, nexttail, advanced), 1, 0, parameters=None ) + @trace_function_calls def _getitem_next_ellipsis( self, tail: tuple[SliceItem, ...], advanced: Index | None ) -> Content: @@ -353,6 +367,7 @@ def _getitem_next_ellipsis( else: return self._getitem_next(slice(None), (Ellipsis, *tail), advanced) + @trace_function_calls def _getitem_next_regular_missing( self, head: IndexedOptionArray, @@ -395,6 +410,7 @@ def _getitem_next_regular_missing( out, indexlength, 1, parameters=self._parameters ) + @trace_function_calls def _getitem_next_missing_jagged( self, head: Content, tail, advanced: Index | None, that: Content ) -> RegularArray: @@ -447,6 +463,7 @@ def _getitem_next_missing_jagged( out, index.length, 1, parameters=self._parameters ) + @trace_function_calls def _getitem_next_missing( self, head: IndexedOptionArray, @@ -508,9 +525,11 @@ def _getitem_next_missing( f"FIXME: unhandled case of SliceMissing with {nextcontent}" ) + @trace_function_calls def __getitem__(self, where): return self._getitem(where) + @trace_function_calls def _getitem(self, where): if is_integer_like(where): return self._getitem_at(ak._slicing.normalize_integer_like(where)) @@ -693,25 +712,31 @@ def _getitem(self, where): + repr(where).replace("\n", "\n ") ) + @trace_function_calls def _is_getitem_at_placeholder(self) -> bool: raise NotImplementedError + @trace_function_calls def _getitem_at(self, where: IndexType): raise NotImplementedError + @trace_function_calls def _getitem_range(self, start: IndexType, stop: IndexType) -> Content: raise NotImplementedError + @trace_function_calls def _getitem_field( self, where: str | SupportsIndex, only_fields: tuple[str, ...] = () ) -> Content: raise NotImplementedError + @trace_function_calls def _getitem_fields( self, where: list[str], only_fields: tuple[str, ...] = () ) -> Content: raise NotImplementedError + @trace_function_calls def _getitem_next( self, head: SliceItem | tuple, @@ -720,6 +745,7 @@ def _getitem_next( ) -> Content: raise NotImplementedError + @trace_function_calls def _getitem_next_jagged( self, slicestarts: Index, @@ -729,9 +755,11 @@ def _getitem_next_jagged( ) -> Content: raise NotImplementedError + @trace_function_calls def _carry(self, carry: Index, allow_lazy: bool) -> Content: raise NotImplementedError + @trace_function_calls def _local_index_axis0(self) -> NumpyArray: localindex = Index64.empty(self.length, self._backend.index_nplike) self._backend.maybe_kernel_error( @@ -744,6 +772,7 @@ def _local_index_axis0(self) -> NumpyArray: localindex.data, parameters=None, backend=self._backend ) + @trace_function_calls def _mergeable_next(self, other: Content, mergebool: bool) -> bool: raise NotImplementedError diff --git a/src/awkward/contents/listarray.py b/src/awkward/contents/listarray.py index a05eeaea55..35b0154024 100644 --- a/src/awkward/contents/listarray.py +++ b/src/awkward/contents/listarray.py @@ -47,6 +47,16 @@ np = NumpyMetadata.instance() +import functools + +def trace_function_calls(func): + @functools.wraps(func) + def wrapper(*args, **kwargs): + print(f"listarray.py: Calling function: {func.__name__}") + result = func(*args, **kwargs) + print(f"listarray.py: Function {func.__name__} returned {result}") + return result + return wrapper @final class ListArray(ListMeta[Content], Content): @@ -298,14 +308,17 @@ def to_RegularArray(self): offsets = self._compact_offsets64(True) return self._broadcast_tooffsets64(offsets).to_RegularArray() + @trace_function_calls def _getitem_nothing(self): return self._content._getitem_range(0, 0) + @trace_function_calls def _is_getitem_at_placeholder(self) -> bool: return isinstance(self._starts, PlaceholderArray) or isinstance( self._stops, PlaceholderArray ) + @trace_function_calls def _getitem_at(self, where: IndexType): if not self._backend.nplike.known_data: self._touch_data(recursive=False) @@ -318,6 +331,7 @@ def _getitem_at(self, where: IndexType): start, stop = self._starts[where], self._stops[where] return self._content._getitem_range(start, stop) + @trace_function_calls def _getitem_range(self, start: IndexType, stop: IndexType) -> Content: if not self._backend.nplike.known_data: self._touch_shape(recursive=False) @@ -330,6 +344,7 @@ def _getitem_range(self, start: IndexType, stop: IndexType) -> Content: parameters=self._parameters, ) + @trace_function_calls def _getitem_field( self, where: str | SupportsIndex, only_fields: tuple[str, ...] = () ) -> Content: @@ -340,6 +355,7 @@ def _getitem_field( parameters=None, ) + @trace_function_calls def _getitem_fields( self, where: list[str | SupportsIndex], only_fields: tuple[str, ...] = () ) -> Content: @@ -350,6 +366,7 @@ def _getitem_fields( parameters=None, ) + @trace_function_calls def _carry(self, carry: Index, allow_lazy: bool) -> Content: assert isinstance(carry, ak.index.Index) @@ -363,6 +380,7 @@ def _carry(self, carry: Index, allow_lazy: bool) -> Content: nextstarts, nextstops, self._content, parameters=self._parameters ) + @trace_function_calls def _compact_offsets64(self, start_at_zero): starts_len = self._starts.length out = ak.index.Index64.empty( @@ -389,6 +407,7 @@ def _compact_offsets64(self, start_at_zero): ) return out + @trace_function_calls def _broadcast_tooffsets64(self, offsets: Index) -> ListOffsetArray: self._touch_data(recursive=False) offsets._touch_data() @@ -443,6 +462,7 @@ def _broadcast_tooffsets64(self, offsets: Index) -> ListOffsetArray: return ListOffsetArray(offsets, nextcontent, parameters=self._parameters) + @trace_function_calls def _getitem_next_jagged( self, slicestarts: Index, slicestops: Index, slicecontent: Content, tail ) -> Content: @@ -698,6 +718,7 @@ def _getitem_next_jagged( f"expected Index/IndexedOptionArray/ListOffsetArray in ListArray._getitem_next_jagged, got {type(slicecontent).__name__}" ) + @trace_function_calls def _getitem_next( self, head: SliceItem | tuple, @@ -1058,9 +1079,11 @@ def _getitem_next( else: raise AssertionError(repr(head)) + @trace_function_calls def _offsets_and_flattened(self, axis: int, depth: int) -> tuple[Index, Content]: return self.to_ListOffsetArray64(True)._offsets_and_flattened(axis, depth) + @trace_function_calls def _mergeable_next(self, other: Content, mergebool: bool) -> bool: # Is the other content is an identity, or a union? if other.is_identity_like or other.is_union: diff --git a/src/awkward/contents/listoffsetarray.py b/src/awkward/contents/listoffsetarray.py index 003467c24b..df38143004 100644 --- a/src/awkward/contents/listoffsetarray.py +++ b/src/awkward/contents/listoffsetarray.py @@ -49,6 +49,16 @@ np = NumpyMetadata.instance() numpy = Numpy.instance() +import functools + +def trace_function_calls(func): + @functools.wraps(func) + def wrapper(*args, **kwargs): + print(f"listoffsetarray.py: Calling function: {func.__name__}") + result = func(*args, **kwargs) + print(f"listoffsetArray.py: Function {func.__name__} returned {result}") + return result + return wrapper @final class ListOffsetArray(ListOffsetMeta[Content], Content): @@ -298,12 +308,15 @@ def to_RegularArray(self): content, size, length, parameters=self._parameters ) + @trace_function_calls def _getitem_nothing(self): return self._content._getitem_range(0, 0) + @trace_function_calls def _is_getitem_at_placeholder(self) -> bool: return isinstance(self._offsets, PlaceholderArray) + @trace_function_calls def _getitem_at(self, where: IndexType): # Wrap `where` by length if not is_unknown_scalar(where) and where < 0: @@ -319,6 +332,7 @@ def _getitem_at(self, where: IndexType): start, stop = self._offsets[where], self._offsets[where + 1] return self._content._getitem_range(start, stop) + @trace_function_calls def _getitem_range(self, start: IndexType, stop: IndexType) -> Content: if not self._backend.nplike.known_data: self._touch_shape(recursive=False) @@ -332,6 +346,7 @@ def _getitem_range(self, start: IndexType, stop: IndexType) -> Content: ) return ListOffsetArray(offsets, self._content, parameters=self._parameters) + @trace_function_calls def _getitem_field( self, where: str | SupportsIndex, only_fields: tuple[str, ...] = () ) -> Content: @@ -341,6 +356,7 @@ def _getitem_field( parameters=None, ) + @trace_function_calls def _getitem_fields( self, where: list[str | SupportsIndex], only_fields: tuple[str, ...] = () ) -> Content: @@ -350,6 +366,7 @@ def _getitem_fields( parameters=None, ) + @trace_function_calls def _carry(self, carry: Index, allow_lazy: bool) -> Content: assert isinstance(carry, ak.index.Index) @@ -363,6 +380,7 @@ def _carry(self, carry: Index, allow_lazy: bool) -> Content: nextstarts, nextstops, self._content, parameters=self._parameters ) + @trace_function_calls def _compact_offsets64(self, start_at_zero: bool) -> Index64: if not start_at_zero or ( self._backend.index_nplike.known_data and self._offsets[0] == 0 @@ -374,6 +392,7 @@ def _compact_offsets64(self, start_at_zero: bool) -> Index64: nplike=self._backend.index_nplike, ) + @trace_function_calls def _broadcast_tooffsets64(self, offsets: Index) -> ListOffsetArray: self._touch_data(recursive=False) offsets._touch_data() @@ -415,6 +434,7 @@ def _broadcast_tooffsets64(self, offsets: Index) -> ListOffsetArray: offsets, next_content[: offsets[-1]], parameters=self._parameters ) + @trace_function_calls def _getitem_next_jagged( self, slicestarts: Index, slicestops: Index, slicecontent: Content, tail ) -> Content: @@ -423,6 +443,7 @@ def _getitem_next_jagged( ) return out._getitem_next_jagged(slicestarts, slicestops, slicecontent, tail) + @trace_function_calls def _getitem_next( self, head: SliceItem | tuple, @@ -700,6 +721,7 @@ def _getitem_next( else: raise AssertionError(repr(head)) + @trace_function_calls def _offsets_and_flattened(self, axis: int, depth: int) -> tuple[Index, Content]: posaxis = maybe_posaxis(self, axis, depth) if posaxis is not None and posaxis + 1 == depth: @@ -776,6 +798,7 @@ def _offsets_and_flattened(self, axis: int, depth: int) -> tuple[Index, Content] ListOffsetArray(tooffsets, flattened, parameters=self._parameters), ) + @trace_function_calls def _mergeable_next(self, other: Content, mergebool: bool) -> bool: # Is the other content is an identity, or a union? if other.is_identity_like or other.is_union: diff --git a/src/awkward/contents/numpyarray.py b/src/awkward/contents/numpyarray.py index 315d9383b7..ef9eb8a2e4 100644 --- a/src/awkward/contents/numpyarray.py +++ b/src/awkward/contents/numpyarray.py @@ -56,6 +56,16 @@ np = NumpyMetadata.instance() numpy = Numpy.instance() +import functools + +def trace_function_calls(func): + @functools.wraps(func) + def wrapper(*args, **kwargs): + print(f"numpyarray.py: Calling function: {func.__name__}") + result = func(*args, **kwargs) + print(f"numpyarray.py: Function {func.__name__} returned {result}") + return result + return wrapper @final class NumpyArray(NumpyMeta, Content): @@ -299,6 +309,7 @@ def maybe_to_NumpyArray(self) -> Self: def __iter__(self): return iter(self._data) + @trace_function_calls def _getitem_nothing(self): tmp = self._data[0:0] return NumpyArray( @@ -307,9 +318,11 @@ def _getitem_nothing(self): backend=self._backend, ) + @trace_function_calls def _is_getitem_at_placeholder(self) -> bool: return isinstance(self._data, PlaceholderArray) + @trace_function_calls def _getitem_at(self, where: IndexType): if not self._backend.nplike.known_data and len(self._data.shape) == 1: self._touch_data(recursive=False) @@ -325,6 +338,7 @@ def _getitem_at(self, where: IndexType): else: return out + @trace_function_calls def _getitem_range(self, start: IndexType, stop: IndexType) -> Content: try: out = self._data[start:stop] @@ -333,11 +347,13 @@ def _getitem_range(self, start: IndexType, stop: IndexType) -> Content: return NumpyArray(out, parameters=self._parameters, backend=self._backend) + @trace_function_calls def _getitem_field( self, where: str | SupportsIndex, only_fields: tuple[str, ...] = () ) -> Content: raise ak._errors.index_error(self, where, "not an array of records") + @trace_function_calls def _getitem_fields( self, where: list[str | SupportsIndex], only_fields: tuple[str, ...] = () ) -> Content: @@ -345,6 +361,7 @@ def _getitem_fields( return self._getitem_range(0, 0) raise ak._errors.index_error(self, where, "not an array of records") + @trace_function_calls def _carry(self, carry: Index, allow_lazy: bool) -> Content: assert isinstance(carry, ak.index.Index) try: @@ -353,6 +370,7 @@ def _carry(self, carry: Index, allow_lazy: bool) -> Content: raise ak._errors.index_error(self, carry.data, str(err)) from err return NumpyArray(nextdata, parameters=self._parameters, backend=self._backend) + @trace_function_calls def _getitem_next_jagged( self, slicestarts: Index, slicestops: Index, slicecontent: Content, tail ) -> Content: @@ -370,6 +388,7 @@ def _getitem_next_jagged( slicestarts, slicestops, slicecontent, tail ) + @trace_function_calls def _getitem_next( self, head: SliceItem | tuple, @@ -440,6 +459,7 @@ def _getitem_next( else: raise AssertionError(repr(head)) + @trace_function_calls def _offsets_and_flattened(self, axis: int, depth: int) -> tuple[Index, Content]: posaxis = maybe_posaxis(self, axis, depth) if posaxis is not None and posaxis + 1 == depth: @@ -451,6 +471,7 @@ def _offsets_and_flattened(self, axis: int, depth: int) -> tuple[Index, Content] else: raise AxisError(f"axis={axis} exceeds the depth of this array ({depth})") + @trace_function_calls def _mergeable_next(self, other: Content, mergebool: bool) -> bool: # Is the other content is an identity, or a union? if other.is_identity_like or other.is_union: diff --git a/tests-cuda/test_3140_cuda_slicing.py b/tests-cuda/test_3140_cuda_slicing.py index 59e2cfcb67..c26f8aea3c 100644 --- a/tests-cuda/test_3140_cuda_slicing.py +++ b/tests-cuda/test_3140_cuda_slicing.py @@ -9,671 +9,696 @@ to_list = ak.operations.to_list -@pytest.fixture(scope="function", autouse=True) +@pytest.fixture(autouse=True) def cleanup_cuda(): yield cp._default_memory_pool.free_all_blocks() cp.cuda.Device().synchronize() -def test_0315_integerindex_null_more(): - f = ak.highlevel.Array([[0, None, 2], None, [3, 4], []], backend="cuda").layout - g1 = ak.highlevel.Array([[1, 2, None], None, [], [None]], backend="cuda").layout - g2 = ak.highlevel.Array([[], None, None, []], backend="cuda").layout - g3 = ak.highlevel.Array([[], [], [], []], backend="cuda").layout - - assert to_list(f[g1]) == [[None, 2, None], None, [], [None]] - assert to_list(f[g2]) == [[], None, None, []] - assert to_list(f[g3]) == [[], None, [], []] - assert f.to_typetracer()[g1].form == f[g1].form - assert f.to_typetracer()[g2].form == f[g2].form - assert f.to_typetracer()[g3].form == f[g3].form - - a = ak.highlevel.Array([[0, 1, 2, None], None], backend="cuda").layout - b = ak.highlevel.Array([[2, 1, None, 3], None], backend="cuda").layout - - assert to_list(a[b]) == [[2, 1, None, None], None] - assert a.to_typetracer()[b].form == a[b].form - - b = ak.highlevel.Array([[2, 1, None, 3], []], backend="cuda").layout - - assert to_list(a[b]) == [[2, 1, None, None], None] - assert a.to_typetracer()[b].form == a[b].form - - b = ak.highlevel.Array([[2, 1, None, 3], [0, 1]], backend="cuda").layout - assert to_list(a[b]) == [[2, 1, None, None], None] - assert a.to_typetracer()[b].form == a[b].form - - -def test_0315_integerindex_null_more_2(): - a = ak.highlevel.Array( - [[[0, 1, 2, None], None], [[3, 4], [5]], None, [[6]]], backend="cuda" - ).layout - b = ak.highlevel.Array( - [[[2, 1, None, 3], [0, 1]], [[0], None], None, [None]], backend="cuda" - ).layout - c = ak.highlevel.Array( - [ - [[False, True, None, False], [False, True]], - [[True, False], None], - None, - [None], - ], - backend="cuda", - ).layout - - assert to_list(a[b]) == [ - [[2, 1, None, None], None], - [[3], None], - None, - [None], - ] - assert to_list(a[c]) == [[[1, None], None], [[4], None], None, [None]] - assert a.to_typetracer()[c].form == a[c].form - - -def test_1405_slicing_untested_cases_list_option_list(): - """Check that non-offset list(option(list indexes correctly""" - content = ak.contents.ListOffsetArray( - ak.index.Index64(np.array([0, 2], dtype=np.int64)), - ak.contents.ListOffsetArray( - ak.index.Index64(np.array([2, 2, 3], dtype=np.int64)), - ak.contents.NumpyArray(np.array([2, 2, 2], dtype=np.int64)), - ), - ) - - index = ak.contents.ListOffsetArray( - ak.index.Index64(np.array([0, 2], dtype=np.int64)), - ak.contents.IndexedOptionArray( - ak.index.Index64(np.array([0, 1], dtype=np.int64)), - ak.contents.ListOffsetArray( - ak.index.Index64(np.array([0, 0, 1], dtype=np.int64)), - ak.contents.NumpyArray(np.array([0], dtype=np.int64)), - ), - ), - ) - - cuda_content = ak.to_backend(content, "cuda") - cuda_index = ak.to_backend(index, "cuda") - - assert cuda_content[cuda_index].to_list() == [[[], [2]]] - - -def test_1405_slicing_untested_cases_list_option_list_offset(): - """Check that offset list(option(list indexes correctly""" - content = ak.contents.ListOffsetArray( - ak.index.Index64(np.array([1, 3], dtype=np.int64)), - ak.contents.ListOffsetArray( - ak.index.Index64(np.array([0, 2, 2, 3], dtype=np.int64)), - ak.contents.NumpyArray(np.array([2, 2, 2], dtype=np.int64)), - ), - ) - - index = ak.contents.ListOffsetArray( - ak.index.Index64(np.array([0, 2], dtype=np.int64)), - ak.contents.IndexedOptionArray( - ak.index.Index64(np.array([0, 1], dtype=np.int64)), - ak.contents.ListOffsetArray( - ak.index.Index64(np.array([0, 0, 1], dtype=np.int64)), - ak.contents.NumpyArray(np.array([0], dtype=np.int64)), - ), - ), - ) - - cuda_content = ak.to_backend(content, "cuda") - cuda_index = ak.to_backend(index, "cuda") - - assert cuda_content[cuda_index].to_list() == [[[], [2]]] - - -def test_1502_getitem_jagged_issue1406(): - array = ak.Array( - ak.contents.ListOffsetArray( - ak.index.Index64(np.array([1, 3], dtype=np.int64)), - ak.contents.ListOffsetArray( - ak.index.Index64(np.array([0, 2, 2, 3], dtype=np.int64)), - ak.contents.NumpyArray(np.array([0, 1, 2], dtype=np.int64)), - ), - ), - check_valid=True, - ) - - index = ak.Array( - ak.contents.ListOffsetArray( - ak.index.Index64(np.array([0, 2], dtype=np.int64)), - ak.contents.IndexedOptionArray( - ak.index.Index64(np.array([0, 1], dtype=np.int64)), - ak.contents.ListOffsetArray( - ak.index.Index64(np.array([0, 0, 1], dtype=np.int64)), - ak.contents.NumpyArray(np.array([0], dtype=np.int64)), - ), - ), - ), - check_valid=True, - ) - - cuda_array = ak.to_backend(array, "cuda") - cuda_index = ak.to_backend(index, "cuda") - - assert to_list(cuda_array[cuda_index]) == [[[], [2]]] - - -def test_1502_getitem_jagged_issue1406_success_start_offset0(): - array = ak.Array( - ak.contents.ListOffsetArray( - ak.index.Index64(np.array([0, 2], dtype=np.int64)), - ak.contents.ListOffsetArray( - ak.index.Index64(np.array([2, 2, 3], dtype=np.int64)), - ak.contents.NumpyArray(np.array([0, 1, 2], dtype=np.int64)), - ), - ), - check_valid=True, - ) - - index = ak.Array( - ak.contents.ListOffsetArray( - ak.index.Index64(np.array([0, 2], dtype=np.int64)), - ak.contents.IndexedOptionArray( - ak.index.Index64(np.array([0, 1], dtype=np.int64)), - ak.contents.ListOffsetArray( - ak.index.Index64(np.array([0, 0, 1], dtype=np.int64)), - ak.contents.NumpyArray(np.array([0], dtype=np.int64)), - ), - ), - ), - check_valid=True, - ) - - cuda_array = ak.to_backend(array, "cuda") - cuda_index = ak.to_backend(index, "cuda") - - assert to_list(cuda_array[cuda_index]) == [[[], [2]]] - - -def test_1502_getitem_jagged_issue1406_success_remove_option_type(): - array = ak.Array( - ak.contents.ListOffsetArray( - ak.index.Index64(np.array([1, 3], dtype=np.int64)), - ak.contents.ListOffsetArray( - ak.index.Index64(np.array([0, 2, 2, 3], dtype=np.int64)), - ak.contents.NumpyArray(np.array([0, 1, 2], dtype=np.int64)), - ), - ), - check_valid=True, - ) - - index = ak.Array( - ak.contents.ListOffsetArray( - ak.index.Index64(np.array([0, 2], dtype=np.int64)), - ak.contents.ListOffsetArray( - ak.index.Index64(np.array([0, 0, 1], dtype=np.int64)), - ak.contents.NumpyArray(np.array([0], dtype=np.int64)), - ), - ), - check_valid=True, - ) - - cuda_array = ak.to_backend(array, "cuda") - cuda_index = ak.to_backend(index, "cuda") - - assert to_list(cuda_array[cuda_index]) == [[[], [2]]] - - -def test_1502_getitem_jagged_issue1406_success_nonempty_list(): - array = ak.Array( - ak.contents.ListOffsetArray( - ak.index.Index64(np.array([1, 3], dtype=np.int64)), - ak.contents.ListOffsetArray( - ak.index.Index64(np.array([0, 1, 2, 3], dtype=np.int64)), - ak.contents.NumpyArray(np.array([0, 1, 2], dtype=np.int64)), - ), - ), - check_valid=True, - ) - - index = ak.Array( - ak.contents.ListOffsetArray( - ak.index.Index64(np.array([0, 2], dtype=np.int64)), - ak.contents.IndexedOptionArray( - ak.index.Index64(np.array([0, 1], dtype=np.int64)), - ak.contents.ListOffsetArray( - ak.index.Index64(np.array([0, 1, 2], dtype=np.int64)), - ak.contents.NumpyArray(np.array([0, 0], dtype=np.int64)), - ), - ), - ), - check_valid=True, - ) - - cuda_array = ak.to_backend(array, "cuda") - cuda_index = ak.to_backend(index, "cuda") - - assert to_list(cuda_array[cuda_index]) == [[[1], [2]]] - - -def test_1904_drop_none_ListArray_and_axis_None(): - index = ak.index.Index64(np.asarray([0, -1, 1, -1, 4, -1, 5])) - content = ak.contents.recordarray.RecordArray( - [ - ak.contents.numpyarray.NumpyArray( - np.array([6.6, 1.1, 2.2, 3.3, 4.4, 5.5, 7.7]) - ) - ], - ["nest"], - ) - indexoptionarray = ak.contents.IndexedOptionArray(index, content) - a = ak.contents.listarray.ListArray( - ak.index.Index(np.array([4, 100, 1], np.int64)), - ak.index.Index(np.array([7, 100, 3, 200], np.int64)), - indexoptionarray, - ) - - cuda_a = ak.to_backend(a, "cuda") - - assert ( - to_list(ak.drop_none(cuda_a)) - == to_list(cuda_a[~ak.is_none(cuda_a, axis=1)]) - == [ - [{"nest": 4.4}, {"nest": 5.5}], - [], - [{"nest": 1.1}], - ] - ) - - -def test_1904_drop_none_ListOffsetArray_IndexedOptionArray_NumpyArray_outoforder(): - index = ak.index.Index64(np.asarray([0, -1, 1, 5, 4, 2, 5])) - content = ak.contents.numpyarray.NumpyArray( - np.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6]) - ) - indexoptionarray = ak.contents.IndexedOptionArray(index, content) - offsets = ak.index.Index64(np.asarray([0, 4, 5, 6])) - listoffset = ak.contents.ListOffsetArray(offsets, indexoptionarray) - - cuda_listoffset = ak.to_backend(listoffset, "cuda") - - assert to_list(cuda_listoffset) == [[0.0, None, 1.1, 5.5], [4.4], [2.2]] - assert ( - to_list(ak.drop_none(cuda_listoffset, axis=1)) - == to_list(cuda_listoffset[~ak.is_none(cuda_listoffset, axis=1)]) - == [[0.0, 1.1, 5.5], [4.4], [2.2]] - ) - assert to_list(ak.drop_none(cuda_listoffset)) == [[0.0, 1.1, 5.5], [4.4], [2.2]] - - -def test_1904_drop_none_from_iter(): - a = ak.Array([[1], [2, None]], backend="cuda") - assert to_list(ak.drop_none(a)) == [[1], [2]] - - a = ak.Array([[2, None]], backend="cuda") - assert to_list(ak.drop_none(a)) == [[2]] - - a = ak.Array([[[None]]], backend="cuda") - assert to_list(ak.drop_none(a)) == [[[]]] - - a = ak.Array([1, 2, None], backend="cuda") - assert to_list(ak.drop_none(a, axis=0)) - - a = ak.Array([[[1, None]], [[3, 4]], [[5, 6]], [[7.8]]], backend="cuda") - assert ( - to_list(ak.drop_none(a, axis=2)) - == to_list(a[~ak.is_none(a, axis=2)]) - == [[[1.0]], [[3.0, 4.0]], [[5.0, 6.0]], [[7.8]]] - ) - - a = ak.Array([[[0]], [[None]], [[1], None], [[2, None]]], backend="cuda") - assert ( - to_list(ak.drop_none(a, axis=1)) - == to_list(a[~ak.is_none(a, axis=1)]) - == [[[0]], [[None]], [[1]], [[2, None]]] - ) - - a = ak.Array( - [[[0]], [None, 34], [[1], None, 31], [[2, [[None]]]], [[[None]]]], - backend="cuda", - ) - assert ( - to_list(ak.drop_none(a, axis=0)) - == to_list(a[~ak.is_none(a, axis=0)]) - == [[[0]], [None, 34], [[1], None, 31], [[2, [[None]]]], [[[None]]]] - ) - - a = ak.Array([[[1, None]], [[3, None]], [[5, 6]], [[7.8]]], backend="cuda") - assert ( - to_list(ak.drop_none(a, axis=2)) - == to_list(a[~ak.is_none(a, axis=2)]) - == [[[1.0]], [[3.0]], [[5.0, 6.0]], [[7.8]]] - ) - - a = ak.Array([[[1, None]], [[None, 4]], [[5, 6]], [[7.8]]], backend="cuda") - assert ( - to_list(ak.drop_none(a, axis=2)) - == to_list(a[~ak.is_none(a, axis=2)]) - == [[[1.0]], [[4.0]], [[5.0, 6.0]], [[7.8]]] - ) - - a = ak.Array([[[1, None]], [[None, None]], [[5, 6]], [[7.8]]], backend="cuda") - assert ( - to_list(ak.drop_none(a, axis=2)) - == to_list(a[~ak.is_none(a, axis=2)]) - == [[[1.0]], [[]], [[5.0, 6.0]], [[7.8]]] - ) - - a = ak.Array([[[1, None]], [[None, None]], [[None, 6]], [[7.8]]], backend="cuda") - assert ( - to_list(ak.drop_none(a, axis=2)) - == to_list(a[~ak.is_none(a, axis=2)]) - == [[[1.0]], [[]], [[6.0]], [[7.8]]] - ) - - a = ak.Array( - [[{"x": [1], "y": [[2]]}], [{"x": [None], "y": [[None]]}], None], backend="cuda" - ) - assert to_list(a) == [ - [{"x": [1], "y": [[2]]}], - [{"x": [None], "y": [[None]]}], - None, - ] - assert to_list(ak.drop_none(a)) == [ - [{"x": [1], "y": [[2]]}], - [{"x": [], "y": [[]]}], - ] - assert ( - to_list(ak.drop_none(a, axis=0)) - == to_list(a[~ak.is_none(a, axis=0)]) - == [[{"x": [1], "y": [[2]]}], [{"x": [None], "y": [[None]]}]] - ) - assert to_list(ak.drop_none(a, axis=1)) == [ - [{"x": [1], "y": [[2]]}], - [{"x": [], "y": [[None]]}], - None, - ] - - -def test_1904_drop_none_List_ByteMaskedArray_NumpyArray(): - a = ak.contents.listarray.ListArray( - ak.index.Index(np.array([1, 3], np.int64)), - ak.index.Index(np.array([3, 4], np.int64)), - ak.contents.bytemaskedarray.ByteMaskedArray( - ak.index.Index(np.array([1, 0, 1, 0, 1], np.int8)), - ak.contents.numpyarray.NumpyArray(np.array([1.1, 2.2, 3.3, 4.4, 5.5, 6.6])), - valid_when=True, - ), - ) - - cuda_a = ak.to_backend(a, "cuda") - - assert to_list(cuda_a) == [[None, 3.3], [None]] - assert to_list(ak.drop_none(cuda_a)) == [[3.3], []] - assert to_list(ak.drop_none(cuda_a, axis=1)) == to_list( - cuda_a[~ak.is_none(cuda_a, axis=1)] - ) - - -def test_1904_drop_none_RegularArray_RecordArray_NumpyArray(): - index = ak.index.Index64(np.asarray([0, -1, 1, 2, 3, 4, -1, 6, 7, 8, -1, 10])) - content = ak.contents.numpyarray.NumpyArray( - np.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9, 10.10]) - ) - indexoptionarray = ak.contents.IndexedOptionArray(index, content) - offsets = ak.index.Index64(np.array([0, 3, 3, 5, 6, 10, 10])) - listoffsetarray = ak.contents.listoffsetarray.ListOffsetArray( - offsets, indexoptionarray - ) - regulararray = ak.contents.regulararray.RegularArray(listoffsetarray, 2) - - cuda_regulararray = ak.to_backend(regulararray, "cuda") - - assert to_list(cuda_regulararray) == [ - [[0.0, None, 1.1], []], - [[2.2, 3.3], [4.4]], - [[None, 6.6, 7.7, 8.8], []], - ] - assert to_list(ak.drop_none(cuda_regulararray, axis=2)) == to_list( - cuda_regulararray[~ak.is_none(cuda_regulararray, axis=2)] - ) - - -def test_1904_drop_none_RecordArray(): - a = ak.Array( - [ - [{"x": [1], "y": [[2]]}], - None, - [None], - [{"x": None, "y": None}], - [{"x": [None], "y": [None]}], - [{"x": [11], "y": [[None]]}], - ], - backend="cuda", - ) - assert to_list(ak.drop_none(a, axis=1)) == to_list(a[~ak.is_none(a, axis=1)]) - assert to_list(ak.drop_none(a, axis=2)) == [ - [{"x": [1], "y": [[2]]}], - None, - [None], - [{"x": None, "y": None}], - [{"x": [], "y": []}], - [{"x": [11], "y": [[None]]}], - ] - - -def test_2246_slice_not_packed(): - index = ak.Array( - ak.contents.ListOffsetArray( - ak.index.Index64([0, 3, 5]), - ak.contents.NumpyArray( - np.array([True, False, False, True, True, False, False], dtype=np.bool_) - ), - ) - ) - array = ak.Array([[0, 1, 2], [3, 4]]) - - cuda_index = ak.to_backend(index, "cuda") - cuda_array = ak.to_backend(array, "cuda") - - result = cuda_array[cuda_index] - assert result.tolist() == [[0], [3, 4]] - - -def test_0127_tomask_operation_ByteMaskedArray_jaggedslice0(): - array = ak.operations.from_iter( - [[0.0, 1.1, 2.2], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9]], highlevel=False - ) - index = ak.index.Index64(np.array([0, 1, 2, 3], dtype=np.int64)) - indexedarray = ak.contents.IndexedOptionArray(index, array) - - cuda_indexedarray = ak.to_backend(indexedarray, "cuda") - - assert to_list(cuda_indexedarray) == [ - [0.0, 1.1, 2.2], - [3.3, 4.4], - [5.5], - [6.6, 7.7, 8.8, 9.9], - ] - assert to_list( - cuda_indexedarray[ - ak.highlevel.Array([[0, -1], [0], [], [1, 1]], backend="cuda") - ] - ) == [ - [0.0, 2.2], - [3.3], - [], - [7.7, 7.7], - ] - - mask = ak.index.Index8(np.array([0, 0, 0, 0], dtype=np.int8)) - maskedarray = ak.contents.ByteMaskedArray(mask, array, valid_when=False) - - cuda_maskedarray = ak.to_backend(maskedarray, "cuda") - - assert to_list(cuda_maskedarray) == [ - [0.0, 1.1, 2.2], - [3.3, 4.4], - [5.5], - [6.6, 7.7, 8.8, 9.9], - ] - assert to_list( - cuda_maskedarray[ak.highlevel.Array([[0, -1], [0], [], [1, 1]], backend="cuda")] - ) == [ - [0.0, 2.2], - [3.3], - [], - [7.7, 7.7], - ] - - -def test_0127_tomask_operation_ByteMaskedArray_jaggedslice1(): - model = ak.highlevel.Array( - [ - [0.0, 1.1, None, 2.2], - [], - [3.3, None, 4.4], - [5.5], - [6.6, 7.7, None, 8.8, 9.9], - ], - backend="cuda", - ) - assert to_list( - model[ - ak.highlevel.Array( - [[3, 2, 1, 1, 0], [], [1], [0, 0], [1, 2]], backend="cuda" - ) - ] - ) == [ - [2.2, None, 1.1, 1.1, 0.0], - [], - [None], - [5.5, 5.5], - [7.7, None], - ] - - content = ak.contents.NumpyArray( - np.array([0.0, 1.1, 999, 2.2, 3.3, 123, 4.4, 5.5, 6.6, 7.7, 321, 8.8, 9.9]) - ) - mask = ak.index.Index8( - np.array([0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0], dtype=np.int8) - ) - maskedarray = ak.contents.ByteMaskedArray(mask, content, valid_when=False) - offsets = ak.index.Index64(np.array([0, 4, 4, 7, 8, 13], dtype=np.int64)) - listarray = ak.highlevel.Array( - ak.contents.ListOffsetArray(offsets, maskedarray), backend="cuda" - ) - - assert to_list(listarray) == to_list(model) - assert to_list( - listarray[ - ak.highlevel.Array( - [[3, 2, 1, 1, 0], [], [1], [0, 0], [1, 2]], backend="cuda" - ) - ] - ) == [[2.2, None, 1.1, 1.1, 0.0], [], [None], [5.5, 5.5], [7.7, None]] - - -def test_0127_tomask_operation_ByteMaskedArray_jaggedslice2(): - model = ak.highlevel.Array( - [ - [[0.0, 1.1, None, 2.2], [], [3.3, None, 4.4]], - [], - [[5.5]], - [[6.6, 7.7, None, 8.8, 9.9]], - ], - backend="cuda", - ) - assert to_list( - model[ - ak.highlevel.Array( - [[[3, 2, 1, 1, 0], [], [1]], [], [[0, 0]], [[1, 2]]], backend="cuda" - ) - ] - ) == [[[2.2, None, 1.1, 1.1, 0.0], [], [None]], [], [[5.5, 5.5]], [[7.7, None]]] - - content = ak.contents.NumpyArray( - np.array([0.0, 1.1, 999, 2.2, 3.3, 123, 4.4, 5.5, 6.6, 7.7, 321, 8.8, 9.9]) - ) - mask = ak.index.Index8( - np.array([0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0], dtype=np.int8) - ) - maskedarray = ak.contents.ByteMaskedArray(mask, content, valid_when=False) - offsets = ak.index.Index64(np.array([0, 4, 4, 7, 8, 13], dtype=np.int64)) - sublistarray = ak.contents.ListOffsetArray(offsets, maskedarray) - offsets2 = ak.index.Index64(np.array([0, 3, 3, 4, 5], dtype=np.int64)) - listarray = ak.highlevel.Array( - ak.contents.ListOffsetArray(offsets2, sublistarray), backend="cuda" - ) - assert to_list(listarray) == to_list(model) - assert to_list( - listarray[ - ak.highlevel.Array( - [[[3, 2, 1, 1, 0], [], [1]], [], [[0, 0]], [[1, 2]]], backend="cuda" - ) - ] - ) == [[[2.2, None, 1.1, 1.1, 0.0], [], [None]], [], [[5.5, 5.5]], [[7.7, None]]] - - -def test_0127_tomask_operation_ByteMaskedArray_jaggedslice3(): - model = ak.highlevel.Array( - [ - [[[0.0, 1.1, None, 2.2], [], [3.3, None, 4.4]], []], - [[[5.5]], [[6.6, 7.7, None, 8.8, 9.9]]], - ], - backend="cuda", - ) - assert to_list( - model[ - ak.highlevel.Array( - [[[[3, 2, 1, 1, 0], [], [1]], []], [[[0, 0]], [[1, 2]]]], backend="cuda" - ) - ] - ) == [[[[2.2, None, 1.1, 1.1, 0.0], [], [None]], []], [[[5.5, 5.5]], [[7.7, None]]]] - - content = ak.contents.NumpyArray( - np.array([0.0, 1.1, 999, 2.2, 3.3, 123, 4.4, 5.5, 6.6, 7.7, 321, 8.8, 9.9]) - ) - mask = ak.index.Index8( - np.array([0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0], dtype=np.int8) - ) - maskedarray = ak.contents.ByteMaskedArray(mask, content, valid_when=False) - offsets = ak.index.Index64(np.array([0, 4, 4, 7, 8, 13], dtype=np.int64)) - subsublistarray = ak.contents.ListOffsetArray(offsets, maskedarray) - offsets2 = ak.index.Index64(np.array([0, 3, 3, 4, 5], dtype=np.int64)) - sublistarray = ak.contents.ListOffsetArray(offsets2, subsublistarray) - offsets3 = ak.index.Index64(np.array([0, 2, 4], dtype=np.int64)) - listarray = ak.highlevel.Array( - ak.contents.ListOffsetArray(offsets3, sublistarray), backend="cuda" - ) - assert to_list(listarray) == to_list(model) - assert to_list( - listarray[ - ak.highlevel.Array( - [[[[3, 2, 1, 1, 0], [], [1]], []], [[[0, 0]], [[1, 2]]]], backend="cuda" - ) - ] - ) == [[[[2.2, None, 1.1, 1.1, 0.0], [], [None]], []], [[[5.5, 5.5]], [[7.7, None]]]] - - -def test_0127_tomask_operation(): - array = ak.highlevel.Array( - [[0.0, 1.1, 2.2], [], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9]], backend="cuda" - ) - mask1 = ak.highlevel.Array([True, True, False, False, True], backend="cuda") - assert to_list(array[mask1]) == [[0.0, 1.1, 2.2], [], [6.6, 7.7, 8.8, 9.9]] - assert to_list(ak.operations.mask(array, mask1)) == [ - [0.0, 1.1, 2.2], - [], - None, - None, - [6.6, 7.7, 8.8, 9.9], - ] - - mask2 = ak.highlevel.Array( - [[False, True, False], [], [True, True], [False], [True, False, False, True]], - backend="cuda", - ) - assert to_list(array[mask2]) == [[1.1], [], [3.3, 4.4], [], [6.6, 9.9]] - assert to_list(ak.operations.mask(array, mask2)) == [ - [None, 1.1, None], - [], - [3.3, 4.4], - [None], - [6.6, None, None, 9.9], - ] +# @pytest.fixture(autouse=True) +# def handle_cuda_errors(): +# try: +# yield +# except cp.cuda.runtime.CUDARuntimeError as e: +# if 'out of memory' in str(e): +# cp.get_default_memory_pool().free_all_blocks() +# pytest.fail("Test failed due to CUDA OOM error.") + + +# def test_0315_integerindex_null_more(): +# f = ak.highlevel.Array([[0, None, 2], None, [3, 4], []], backend="cuda").layout +# g1 = ak.highlevel.Array([[1, 2, None], None, [], [None]], backend="cuda").layout +# g2 = ak.highlevel.Array([[], None, None, []], backend="cuda").layout +# g3 = ak.highlevel.Array([[], [], [], []], backend="cuda").layout + +# assert to_list(f[g1]) == [[None, 2, None], None, [], [None]] +# assert to_list(f[g2]) == [[], None, None, []] +# assert to_list(f[g3]) == [[], None, [], []] +# assert f.to_typetracer()[g1].form == f[g1].form +# assert f.to_typetracer()[g2].form == f[g2].form +# assert f.to_typetracer()[g3].form == f[g3].form + +# a = ak.highlevel.Array([[0, 1, 2, None], None], backend="cuda").layout +# b = ak.highlevel.Array([[2, 1, None, 3], None], backend="cuda").layout + +# assert to_list(a[b]) == [[2, 1, None, None], None] +# assert a.to_typetracer()[b].form == a[b].form + +# b = ak.highlevel.Array([[2, 1, None, 3], []], backend="cuda").layout + +# assert to_list(a[b]) == [[2, 1, None, None], None] +# assert a.to_typetracer()[b].form == a[b].form + +# b = ak.highlevel.Array([[2, 1, None, 3], [0, 1]], backend="cuda").layout +# assert to_list(a[b]) == [[2, 1, None, None], None] +# assert a.to_typetracer()[b].form == a[b].form + + +# def test_0315_integerindex_null_more_2(): +# a = ak.highlevel.Array( +# [[[0, 1, 2, None], None], [[3, 4], [5]], None, [[6]]], backend="cuda" +# ).layout +# b = ak.highlevel.Array( +# [[[2, 1, None, 3], [0, 1]], [[0], None], None, [None]], backend="cuda" +# ).layout +# c = ak.highlevel.Array( +# [ +# [[False, True, None, False], [False, True]], +# [[True, False], None], +# None, +# [None], +# ], +# backend="cuda", +# ).layout + +# assert to_list(a[b]) == [ +# [[2, 1, None, None], None], +# [[3], None], +# None, +# [None], +# ] +# assert to_list(a[c]) == [[[1, None], None], [[4], None], None, [None]] +# assert a.to_typetracer()[c].form == a[c].form + + +# def test_1405_slicing_untested_cases_list_option_list(): +# """Check that non-offset list(option(list indexes correctly""" +# content = ak.contents.ListOffsetArray( +# ak.index.Index64(np.array([0, 2], dtype=np.int64)), +# ak.contents.ListOffsetArray( +# ak.index.Index64(np.array([2, 2, 3], dtype=np.int64)), +# ak.contents.NumpyArray(np.array([2, 2, 2], dtype=np.int64)), +# ), +# ) + +# index = ak.contents.ListOffsetArray( +# ak.index.Index64(np.array([0, 2], dtype=np.int64)), +# ak.contents.IndexedOptionArray( +# ak.index.Index64(np.array([0, 1], dtype=np.int64)), +# ak.contents.ListOffsetArray( +# ak.index.Index64(np.array([0, 0, 1], dtype=np.int64)), +# ak.contents.NumpyArray(np.array([0], dtype=np.int64)), +# ), +# ), +# ) + +# cuda_content = ak.to_backend(content, "cuda") +# cuda_index = ak.to_backend(index, "cuda") + +# assert cuda_content[cuda_index].to_list() == [[[], [2]]] + + +# def test_1405_slicing_untested_cases_list_option_list_offset(): +# """Check that offset list(option(list indexes correctly""" +# content = ak.contents.ListOffsetArray( +# ak.index.Index64(np.array([1, 3], dtype=np.int64)), +# ak.contents.ListOffsetArray( +# ak.index.Index64(np.array([0, 2, 2, 3], dtype=np.int64)), +# ak.contents.NumpyArray(np.array([2, 2, 2], dtype=np.int64)), +# ), +# ) + +# index = ak.contents.ListOffsetArray( +# ak.index.Index64(np.array([0, 2], dtype=np.int64)), +# ak.contents.IndexedOptionArray( +# ak.index.Index64(np.array([0, 1], dtype=np.int64)), +# ak.contents.ListOffsetArray( +# ak.index.Index64(np.array([0, 0, 1], dtype=np.int64)), +# ak.contents.NumpyArray(np.array([0], dtype=np.int64)), +# ), +# ), +# ) + +# cuda_content = ak.to_backend(content, "cuda") +# cuda_index = ak.to_backend(index, "cuda") + +# assert cuda_content[cuda_index].to_list() == [[[], [2]]] + + +# def test_1502_getitem_jagged_issue1406(): +# array = ak.Array( +# ak.contents.ListOffsetArray( +# ak.index.Index64(np.array([1, 3], dtype=np.int64)), +# ak.contents.ListOffsetArray( +# ak.index.Index64(np.array([0, 2, 2, 3], dtype=np.int64)), +# ak.contents.NumpyArray(np.array([0, 1, 2], dtype=np.int64)), +# ), +# ), +# check_valid=True, +# ) + +# index = ak.Array( +# ak.contents.ListOffsetArray( +# ak.index.Index64(np.array([0, 2], dtype=np.int64)), +# ak.contents.IndexedOptionArray( +# ak.index.Index64(np.array([0, 1], dtype=np.int64)), +# ak.contents.ListOffsetArray( +# ak.index.Index64(np.array([0, 0, 1], dtype=np.int64)), +# ak.contents.NumpyArray(np.array([0], dtype=np.int64)), +# ), +# ), +# ), +# check_valid=True, +# ) + +# cuda_array = ak.to_backend(array, "cuda") +# cuda_index = ak.to_backend(index, "cuda") + +# assert to_list(cuda_array[cuda_index]) == [[[], [2]]] + + +# def test_1502_getitem_jagged_issue1406_success_start_offset0(): +# array = ak.Array( +# ak.contents.ListOffsetArray( +# ak.index.Index64(np.array([0, 2], dtype=np.int64)), +# ak.contents.ListOffsetArray( +# ak.index.Index64(np.array([2, 2, 3], dtype=np.int64)), +# ak.contents.NumpyArray(np.array([0, 1, 2], dtype=np.int64)), +# ), +# ), +# check_valid=True, +# ) + +# index = ak.Array( +# ak.contents.ListOffsetArray( +# ak.index.Index64(np.array([0, 2], dtype=np.int64)), +# ak.contents.IndexedOptionArray( +# ak.index.Index64(np.array([0, 1], dtype=np.int64)), +# ak.contents.ListOffsetArray( +# ak.index.Index64(np.array([0, 0, 1], dtype=np.int64)), +# ak.contents.NumpyArray(np.array([0], dtype=np.int64)), +# ), +# ), +# ), +# check_valid=True, +# ) + +# cuda_array = ak.to_backend(array, "cuda") +# cuda_index = ak.to_backend(index, "cuda") + +# assert to_list(cuda_array[cuda_index]) == [[[], [2]]] + + +# def test_1502_getitem_jagged_issue1406_success_remove_option_type(): +# array = ak.Array( +# ak.contents.ListOffsetArray( +# ak.index.Index64(np.array([1, 3], dtype=np.int64)), +# ak.contents.ListOffsetArray( +# ak.index.Index64(np.array([0, 2, 2, 3], dtype=np.int64)), +# ak.contents.NumpyArray(np.array([0, 1, 2], dtype=np.int64)), +# ), +# ), +# check_valid=True, +# ) + +# index = ak.Array( +# ak.contents.ListOffsetArray( +# ak.index.Index64(np.array([0, 2], dtype=np.int64)), +# ak.contents.ListOffsetArray( +# ak.index.Index64(np.array([0, 0, 1], dtype=np.int64)), +# ak.contents.NumpyArray(np.array([0], dtype=np.int64)), +# ), +# ), +# check_valid=True, +# ) + +# cuda_array = ak.to_backend(array, "cuda") +# cuda_index = ak.to_backend(index, "cuda") + +# assert to_list(cuda_array[cuda_index]) == [[[], [2]]] + + +# def test_1502_getitem_jagged_issue1406_success_nonempty_list(): +# array = ak.Array( +# ak.contents.ListOffsetArray( +# ak.index.Index64(np.array([1, 3], dtype=np.int64)), +# ak.contents.ListOffsetArray( +# ak.index.Index64(np.array([0, 1, 2, 3], dtype=np.int64)), +# ak.contents.NumpyArray(np.array([0, 1, 2], dtype=np.int64)), +# ), +# ), +# check_valid=True, +# ) + +# index = ak.Array( +# ak.contents.ListOffsetArray( +# ak.index.Index64(np.array([0, 2], dtype=np.int64)), +# ak.contents.IndexedOptionArray( +# ak.index.Index64(np.array([0, 1], dtype=np.int64)), +# ak.contents.ListOffsetArray( +# ak.index.Index64(np.array([0, 1, 2], dtype=np.int64)), +# ak.contents.NumpyArray(np.array([0, 0], dtype=np.int64)), +# ), +# ), +# ), +# check_valid=True, +# ) + +# cuda_array = ak.to_backend(array, "cuda") +# cuda_index = ak.to_backend(index, "cuda") + +# assert to_list(cuda_array[cuda_index]) == [[[1], [2]]] + + +# def test_1904_drop_none_ListArray_and_axis_None(): +# index = ak.index.Index64(np.asarray([0, -1, 1, -1, 4, -1, 5])) +# content = ak.contents.recordarray.RecordArray( +# [ +# ak.contents.numpyarray.NumpyArray( +# np.array([6.6, 1.1, 2.2, 3.3, 4.4, 5.5, 7.7]) +# ) +# ], +# ["nest"], +# ) +# indexoptionarray = ak.contents.IndexedOptionArray(index, content) +# a = ak.contents.listarray.ListArray( +# ak.index.Index(np.array([4, 100, 1], np.int64)), +# ak.index.Index(np.array([7, 100, 3, 200], np.int64)), +# indexoptionarray, +# ) + +# cuda_a = ak.to_backend(a, "cuda") + +# assert ( +# to_list(ak.drop_none(cuda_a)) +# == to_list(cuda_a[~ak.is_none(cuda_a, axis=1)]) +# == [ +# [{"nest": 4.4}, {"nest": 5.5}], +# [], +# [{"nest": 1.1}], +# ] +# ) + + +# def test_1904_drop_none_ListOffsetArray_IndexedOptionArray_NumpyArray_outoforder(): +# index = ak.index.Index64(np.asarray([0, -1, 1, 5, 4, 2, 5])) +# content = ak.contents.numpyarray.NumpyArray( +# np.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6]) +# ) +# indexoptionarray = ak.contents.IndexedOptionArray(index, content) +# offsets = ak.index.Index64(np.asarray([0, 4, 5, 6])) +# listoffset = ak.contents.ListOffsetArray(offsets, indexoptionarray) + +# cuda_listoffset = ak.to_backend(listoffset, "cuda") + +# assert to_list(cuda_listoffset) == [[0.0, None, 1.1, 5.5], [4.4], [2.2]] +# assert ( +# to_list(ak.drop_none(cuda_listoffset, axis=1)) +# == to_list(cuda_listoffset[~ak.is_none(cuda_listoffset, axis=1)]) +# == [[0.0, 1.1, 5.5], [4.4], [2.2]] +# ) +# assert to_list(ak.drop_none(cuda_listoffset)) == [[0.0, 1.1, 5.5], [4.4], [2.2]] + + +# def test_1904_drop_none_from_iter(): +# a = ak.Array([[1], [2, None]], backend="cuda") +# assert to_list(ak.drop_none(a)) == [[1], [2]] + +# a = ak.Array([[2, None]], backend="cuda") +# assert to_list(ak.drop_none(a)) == [[2]] + +# a = ak.Array([[[None]]], backend="cuda") +# assert to_list(ak.drop_none(a)) == [[[]]] + +# a = ak.Array([1, 2, None], backend="cuda") +# assert to_list(ak.drop_none(a, axis=0)) + +# a = ak.Array([[[1, None]], [[3, 4]], [[5, 6]], [[7.8]]], backend="cuda") +# assert ( +# to_list(ak.drop_none(a, axis=2)) +# == to_list(a[~ak.is_none(a, axis=2)]) +# == [[[1.0]], [[3.0, 4.0]], [[5.0, 6.0]], [[7.8]]] +# ) + +# a = ak.Array([[[0]], [[None]], [[1], None], [[2, None]]], backend="cuda") +# assert ( +# to_list(ak.drop_none(a, axis=1)) +# == to_list(a[~ak.is_none(a, axis=1)]) +# == [[[0]], [[None]], [[1]], [[2, None]]] +# ) + +# a = ak.Array( +# [[[0]], [None, 34], [[1], None, 31], [[2, [[None]]]], [[[None]]]], +# backend="cuda", +# ) +# assert ( +# to_list(ak.drop_none(a, axis=0)) +# == to_list(a[~ak.is_none(a, axis=0)]) +# == [[[0]], [None, 34], [[1], None, 31], [[2, [[None]]]], [[[None]]]] +# ) + +# a = ak.Array([[[1, None]], [[3, None]], [[5, 6]], [[7.8]]], backend="cuda") +# assert ( +# to_list(ak.drop_none(a, axis=2)) +# == to_list(a[~ak.is_none(a, axis=2)]) +# == [[[1.0]], [[3.0]], [[5.0, 6.0]], [[7.8]]] +# ) + +# a = ak.Array([[[1, None]], [[None, 4]], [[5, 6]], [[7.8]]], backend="cuda") +# assert ( +# to_list(ak.drop_none(a, axis=2)) +# == to_list(a[~ak.is_none(a, axis=2)]) +# == [[[1.0]], [[4.0]], [[5.0, 6.0]], [[7.8]]] +# ) + +# a = ak.Array([[[1, None]], [[None, None]], [[5, 6]], [[7.8]]], backend="cuda") +# assert ( +# to_list(ak.drop_none(a, axis=2)) +# == to_list(a[~ak.is_none(a, axis=2)]) +# == [[[1.0]], [[]], [[5.0, 6.0]], [[7.8]]] +# ) + +# a = ak.Array([[[1, None]], [[None, None]], [[None, 6]], [[7.8]]], backend="cuda") +# assert ( +# to_list(ak.drop_none(a, axis=2)) +# == to_list(a[~ak.is_none(a, axis=2)]) +# == [[[1.0]], [[]], [[6.0]], [[7.8]]] +# ) + +# a = ak.Array( +# [[{"x": [1], "y": [[2]]}], [{"x": [None], "y": [[None]]}], None], backend="cuda" +# ) +# assert to_list(a) == [ +# [{"x": [1], "y": [[2]]}], +# [{"x": [None], "y": [[None]]}], +# None, +# ] +# assert to_list(ak.drop_none(a)) == [ +# [{"x": [1], "y": [[2]]}], +# [{"x": [], "y": [[]]}], +# ] +# assert ( +# to_list(ak.drop_none(a, axis=0)) +# == to_list(a[~ak.is_none(a, axis=0)]) +# == [[{"x": [1], "y": [[2]]}], [{"x": [None], "y": [[None]]}]] +# ) +# assert to_list(ak.drop_none(a, axis=1)) == [ +# [{"x": [1], "y": [[2]]}], +# [{"x": [], "y": [[None]]}], +# None, +# ] + + +# def test_1904_drop_none_List_ByteMaskedArray_NumpyArray(): +# a = ak.contents.listarray.ListArray( +# ak.index.Index(np.array([1, 3], np.int64)), +# ak.index.Index(np.array([3, 4], np.int64)), +# ak.contents.bytemaskedarray.ByteMaskedArray( +# ak.index.Index(np.array([1, 0, 1, 0, 1], np.int8)), +# ak.contents.numpyarray.NumpyArray(np.array([1.1, 2.2, 3.3, 4.4, 5.5, 6.6])), +# valid_when=True, +# ), +# ) + +# cuda_a = ak.to_backend(a, "cuda") + +# assert to_list(cuda_a) == [[None, 3.3], [None]] +# assert to_list(ak.drop_none(cuda_a)) == [[3.3], []] +# assert to_list(ak.drop_none(cuda_a, axis=1)) == to_list( +# cuda_a[~ak.is_none(cuda_a, axis=1)] +# ) + + +# def test_1904_drop_none_RegularArray_RecordArray_NumpyArray(): +# index = ak.index.Index64(np.asarray([0, -1, 1, 2, 3, 4, -1, 6, 7, 8, -1, 10])) +# content = ak.contents.numpyarray.NumpyArray( +# np.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9, 10.10]) +# ) +# indexoptionarray = ak.contents.IndexedOptionArray(index, content) +# offsets = ak.index.Index64(np.array([0, 3, 3, 5, 6, 10, 10])) +# listoffsetarray = ak.contents.listoffsetarray.ListOffsetArray( +# offsets, indexoptionarray +# ) +# regulararray = ak.contents.regulararray.RegularArray(listoffsetarray, 2) + +# cuda_regulararray = ak.to_backend(regulararray, "cuda") + +# assert to_list(cuda_regulararray) == [ +# [[0.0, None, 1.1], []], +# [[2.2, 3.3], [4.4]], +# [[None, 6.6, 7.7, 8.8], []], +# ] +# assert to_list(ak.drop_none(cuda_regulararray, axis=2)) == to_list( +# cuda_regulararray[~ak.is_none(cuda_regulararray, axis=2)] +# ) + + +# def test_1904_drop_none_RecordArray(): +# a = ak.Array( +# [ +# [{"x": [1], "y": [[2]]}], +# None, +# [None], +# [{"x": None, "y": None}], +# [{"x": [None], "y": [None]}], +# [{"x": [11], "y": [[None]]}], +# ], +# backend="cuda", +# ) +# assert to_list(ak.drop_none(a, axis=1)) == to_list(a[~ak.is_none(a, axis=1)]) +# assert to_list(ak.drop_none(a, axis=2)) == [ +# [{"x": [1], "y": [[2]]}], +# None, +# [None], +# [{"x": None, "y": None}], +# [{"x": [], "y": []}], +# [{"x": [11], "y": [[None]]}], +# ] + +def test_simple_slice_cpu(): + arrg = ak.Array([[1, 2, 3], [0], [4, 5]]) + out = arrg[:, 0] + expected = [1, 0, 4] + result = out.tolist() + cp.testing.assert_array_list_equal(result, expected, err_msg=f"Slice of [[1, 2, 3], [0], [4, 5]] should be {expected}, but got {result}") + + +def test_simple_slice_gpu(): + arrg = ak.Array([[1, 2, 3], [0], [4, 5]], backend="cuda") + out = arrg[:, 0] + expected = [1, 0, 4] + result = out.tolist() + cp.testing.assert_array_list_equal(result, expected, err_msg=f"Slice of [[1, 2, 3], [0], [4, 5]] should be {expected}, but got {result}") + + +# def test_2246_slice_not_packed(): +# index = ak.Array( +# ak.contents.ListOffsetArray( +# ak.index.Index64([0, 3, 5]), +# ak.contents.NumpyArray( +# np.array([True, False, False, True, True, False, False], dtype=np.bool_) +# ), +# ) +# ) +# array = ak.Array([[0, 1, 2], [3, 4]]) + +# cuda_index = ak.to_backend(index, "cuda") +# cuda_array = ak.to_backend(array, "cuda") + +# result = cuda_array[cuda_index] +# assert result.tolist() == [[0], [3, 4]] + + +# def test_0127_tomask_operation_ByteMaskedArray_jaggedslice0(): +# array = ak.operations.from_iter( +# [[0.0, 1.1, 2.2], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9]], highlevel=False +# ) +# index = ak.index.Index64(np.array([0, 1, 2, 3], dtype=np.int64)) +# indexedarray = ak.contents.IndexedOptionArray(index, array) + +# cuda_indexedarray = ak.to_backend(indexedarray, "cuda") + +# assert to_list(cuda_indexedarray) == [ +# [0.0, 1.1, 2.2], +# [3.3, 4.4], +# [5.5], +# [6.6, 7.7, 8.8, 9.9], +# ] +# assert to_list( +# cuda_indexedarray[ +# ak.highlevel.Array([[0, -1], [0], [], [1, 1]], backend="cuda") +# ] +# ) == [ +# [0.0, 2.2], +# [3.3], +# [], +# [7.7, 7.7], +# ] + +# mask = ak.index.Index8(np.array([0, 0, 0, 0], dtype=np.int8)) +# maskedarray = ak.contents.ByteMaskedArray(mask, array, valid_when=False) + +# cuda_maskedarray = ak.to_backend(maskedarray, "cuda") + +# assert to_list(cuda_maskedarray) == [ +# [0.0, 1.1, 2.2], +# [3.3, 4.4], +# [5.5], +# [6.6, 7.7, 8.8, 9.9], +# ] +# assert to_list( +# cuda_maskedarray[ak.highlevel.Array([[0, -1], [0], [], [1, 1]], backend="cuda")] +# ) == [ +# [0.0, 2.2], +# [3.3], +# [], +# [7.7, 7.7], +# ] + + +# def test_0127_tomask_operation_ByteMaskedArray_jaggedslice1(): +# model = ak.highlevel.Array( +# [ +# [0.0, 1.1, None, 2.2], +# [], +# [3.3, None, 4.4], +# [5.5], +# [6.6, 7.7, None, 8.8, 9.9], +# ], +# backend="cuda", +# ) +# assert to_list( +# model[ +# ak.highlevel.Array( +# [[3, 2, 1, 1, 0], [], [1], [0, 0], [1, 2]], backend="cuda" +# ) +# ] +# ) == [ +# [2.2, None, 1.1, 1.1, 0.0], +# [], +# [None], +# [5.5, 5.5], +# [7.7, None], +# ] + +# content = ak.contents.NumpyArray( +# np.array([0.0, 1.1, 999, 2.2, 3.3, 123, 4.4, 5.5, 6.6, 7.7, 321, 8.8, 9.9]) +# ) +# mask = ak.index.Index8( +# np.array([0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0], dtype=np.int8) +# ) +# maskedarray = ak.contents.ByteMaskedArray(mask, content, valid_when=False) +# offsets = ak.index.Index64(np.array([0, 4, 4, 7, 8, 13], dtype=np.int64)) +# listarray = ak.highlevel.Array( +# ak.contents.ListOffsetArray(offsets, maskedarray), backend="cuda" +# ) + +# assert to_list(listarray) == to_list(model) +# assert to_list( +# listarray[ +# ak.highlevel.Array( +# [[3, 2, 1, 1, 0], [], [1], [0, 0], [1, 2]], backend="cuda" +# ) +# ] +# ) == [[2.2, None, 1.1, 1.1, 0.0], [], [None], [5.5, 5.5], [7.7, None]] + + +# def test_0127_tomask_operation_ByteMaskedArray_jaggedslice2(): +# model = ak.highlevel.Array( +# [ +# [[0.0, 1.1, None, 2.2], [], [3.3, None, 4.4]], +# [], +# [[5.5]], +# [[6.6, 7.7, None, 8.8, 9.9]], +# ], +# backend="cuda", +# ) +# assert to_list( +# model[ +# ak.highlevel.Array( +# [[[3, 2, 1, 1, 0], [], [1]], [], [[0, 0]], [[1, 2]]], backend="cuda" +# ) +# ] +# ) == [[[2.2, None, 1.1, 1.1, 0.0], [], [None]], [], [[5.5, 5.5]], [[7.7, None]]] + +# content = ak.contents.NumpyArray( +# np.array([0.0, 1.1, 999, 2.2, 3.3, 123, 4.4, 5.5, 6.6, 7.7, 321, 8.8, 9.9]) +# ) +# mask = ak.index.Index8( +# np.array([0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0], dtype=np.int8) +# ) +# maskedarray = ak.contents.ByteMaskedArray(mask, content, valid_when=False) +# offsets = ak.index.Index64(np.array([0, 4, 4, 7, 8, 13], dtype=np.int64)) +# sublistarray = ak.contents.ListOffsetArray(offsets, maskedarray) +# offsets2 = ak.index.Index64(np.array([0, 3, 3, 4, 5], dtype=np.int64)) +# listarray = ak.highlevel.Array( +# ak.contents.ListOffsetArray(offsets2, sublistarray), backend="cuda" +# ) +# assert to_list(listarray) == to_list(model) +# assert to_list( +# listarray[ +# ak.highlevel.Array( +# [[[3, 2, 1, 1, 0], [], [1]], [], [[0, 0]], [[1, 2]]], backend="cuda" +# ) +# ] +# ) == [[[2.2, None, 1.1, 1.1, 0.0], [], [None]], [], [[5.5, 5.5]], [[7.7, None]]] + + +# def test_0127_tomask_operation_ByteMaskedArray_jaggedslice3(): +# model = ak.highlevel.Array( +# [ +# [[[0.0, 1.1, None, 2.2], [], [3.3, None, 4.4]], []], +# [[[5.5]], [[6.6, 7.7, None, 8.8, 9.9]]], +# ], +# backend="cuda", +# ) +# assert to_list( +# model[ +# ak.highlevel.Array( +# [[[[3, 2, 1, 1, 0], [], [1]], []], [[[0, 0]], [[1, 2]]]], backend="cuda" +# ) +# ] +# ) == [[[[2.2, None, 1.1, 1.1, 0.0], [], [None]], []], [[[5.5, 5.5]], [[7.7, None]]]] + +# content = ak.contents.NumpyArray( +# np.array([0.0, 1.1, 999, 2.2, 3.3, 123, 4.4, 5.5, 6.6, 7.7, 321, 8.8, 9.9]) +# ) +# mask = ak.index.Index8( +# np.array([0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0], dtype=np.int8) +# ) +# maskedarray = ak.contents.ByteMaskedArray(mask, content, valid_when=False) +# offsets = ak.index.Index64(np.array([0, 4, 4, 7, 8, 13], dtype=np.int64)) +# subsublistarray = ak.contents.ListOffsetArray(offsets, maskedarray) +# offsets2 = ak.index.Index64(np.array([0, 3, 3, 4, 5], dtype=np.int64)) +# sublistarray = ak.contents.ListOffsetArray(offsets2, subsublistarray) +# offsets3 = ak.index.Index64(np.array([0, 2, 4], dtype=np.int64)) +# listarray = ak.highlevel.Array( +# ak.contents.ListOffsetArray(offsets3, sublistarray), backend="cuda" +# ) +# assert to_list(listarray) == to_list(model) +# assert to_list( +# listarray[ +# ak.highlevel.Array( +# [[[[3, 2, 1, 1, 0], [], [1]], []], [[[0, 0]], [[1, 2]]]], backend="cuda" +# ) +# ] +# ) == [[[[2.2, None, 1.1, 1.1, 0.0], [], [None]], []], [[[5.5, 5.5]], [[7.7, None]]]] + + +# def test_0127_tomask_operation(): +# array = ak.highlevel.Array( +# [[0.0, 1.1, 2.2], [], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9]], backend="cuda" +# ) +# mask1 = ak.highlevel.Array([True, True, False, False, True], backend="cuda") +# assert to_list(array[mask1]) == [[0.0, 1.1, 2.2], [], [6.6, 7.7, 8.8, 9.9]] +# assert to_list(ak.operations.mask(array, mask1)) == [ +# [0.0, 1.1, 2.2], +# [], +# None, +# None, +# [6.6, 7.7, 8.8, 9.9], +# ] + +# mask2 = ak.highlevel.Array( +# [[False, True, False], [], [True, True], [False], [True, False, False, True]], +# backend="cuda", +# ) +# assert to_list(array[mask2]) == [[1.1], [], [3.3, 4.4], [], [6.6, 9.9]] +# assert to_list(ak.operations.mask(array, mask2)) == [ +# [None, 1.1, None], +# [], +# [3.3, 4.4], +# [None], +# [6.6, None, None, 9.9], +# ] From 01fd17b58504552a471ee49afb24feda7e339cf5 Mon Sep 17 00:00:00 2001 From: Ianna Osborne Date: Tue, 17 Sep 2024 17:35:41 +0200 Subject: [PATCH 2/2] debug slicing --- src/awkward/_slicing.py | 4 ++++ src/awkward/contents/content.py | 19 +++++++++++++++++++ src/awkward/contents/listoffsetarray.py | 3 +++ src/awkward/contents/numpyarray.py | 5 ++++- tests-cuda/test_3140_cuda_slicing.py | 11 +++++++---- 5 files changed, 37 insertions(+), 5 deletions(-) diff --git a/src/awkward/_slicing.py b/src/awkward/_slicing.py index f8193951be..470941f393 100644 --- a/src/awkward/_slicing.py +++ b/src/awkward/_slicing.py @@ -214,9 +214,11 @@ def normalise_item(item, backend: Backend) -> SliceItem: """ # Basic indices if is_integer_like(item): + print(" normalize_item::integer_like", item) return normalize_integer_like(item) elif isinstance(item, slice): + print(" normalize_item::slice", item) return normalize_slice(item, nplike=backend.index_nplike) elif isinstance(item, str): @@ -317,6 +319,8 @@ def normalise_item(item, backend: Backend) -> SliceItem: @trace_function_calls def normalise_items(where: Sequence, backend: Backend) -> list: # First prepare items for broadcasting into like-types + for x in where: + print(" normalise_items", x, where) return [normalise_item(x, backend=backend) for x in where] @trace_function_calls diff --git a/src/awkward/contents/content.py b/src/awkward/contents/content.py index 5d71f5064e..585228078f 100644 --- a/src/awkward/contents/content.py +++ b/src/awkward/contents/content.py @@ -527,44 +527,57 @@ def _getitem_next_missing( @trace_function_calls def __getitem__(self, where): + print(" content::__getitem__ of ", self, self.backend) return self._getitem(where) @trace_function_calls def _getitem(self, where): if is_integer_like(where): + print(" in content::_getitem: is integer like!", where) return self._getitem_at(ak._slicing.normalize_integer_like(where)) elif isinstance(where, slice) and where.step is None: + print(" content::slice with step???") # Ensure that start, stop are non-negative! start, stop, _, _ = self._backend.index_nplike.derive_slice_for_length( normalize_slice(where, nplike=self._backend.index_nplike), self.length ) + print(" >>>start, stop", start, stop) return self._getitem_range(start, stop) elif isinstance(where, slice): + print(" content::slice???") return self._getitem((where,)) elif isinstance(where, str): + print(" content::str") return self._getitem_field(where) elif where is np.newaxis: + print(" content::axis") return self._getitem((where,)) elif where is Ellipsis: + print(" content::Ellipsis") return self._getitem((where,)) elif isinstance(where, tuple): + print(" content::tuple", where) if len(where) == 0: + print(" len 0") return self # Backend may change if index contains typetracers backend = backend_of(self, *where, coerce_to_common=True) this = self.to_backend(backend) + print(" this", this) # Normalise valid indices onto well-defined basis items = ak._slicing.normalise_items(where, backend) + print(" items", items) # Prepare items for advanced indexing (e.g. via broadcasting) nextwhere = ak._slicing.prepare_advanced_indexing(items, backend) + print(" nextwhere", nextwhere) next = ak.contents.RegularArray( this, @@ -572,8 +585,11 @@ def _getitem(self, where): 1, parameters=None, ) + print(" next", next) + out = next._getitem_next(nextwhere[0], nextwhere[1:], None) + print(" content::tuple out", out) if out.length is not unknown_length and out.length == 0: return out._getitem_nothing() @@ -581,6 +597,7 @@ def _getitem(self, where): return out._getitem_at(0) elif isinstance(where, ak.highlevel.Array): + print(" content::Array") return self._getitem(where.layout) # Convert between nplikes of different backends @@ -588,10 +605,12 @@ def _getitem(self, where): isinstance(where, ak.contents.Content) and where.backend is not self._backend ): + print(" content::backends") backend = backend_of(self, where, coerce_to_common=True) return self.to_backend(backend)._getitem(where.to_backend(backend)) elif isinstance(where, ak.contents.NumpyArray): + print(" content::NumpyArray") data_as_index = to_nplike( where.data, self._backend.index_nplike, diff --git a/src/awkward/contents/listoffsetarray.py b/src/awkward/contents/listoffsetarray.py index df38143004..02ef8783aa 100644 --- a/src/awkward/contents/listoffsetarray.py +++ b/src/awkward/contents/listoffsetarray.py @@ -318,6 +318,7 @@ def _is_getitem_at_placeholder(self) -> bool: @trace_function_calls def _getitem_at(self, where: IndexType): + print(" listoffsetarray::_getitem_at", self, self.backend, where) # Wrap `where` by length if not is_unknown_scalar(where) and where < 0: length_index = self._backend.index_nplike.shape_item_as_index(self.length) @@ -330,6 +331,8 @@ def _getitem_at(self, where: IndexType): ): raise ak._errors.index_error(self, where) start, stop = self._offsets[where], self._offsets[where + 1] + print(" listoffsetarray::_getitem_at start stop", start, stop) + return self._content._getitem_range(start, stop) @trace_function_calls diff --git a/src/awkward/contents/numpyarray.py b/src/awkward/contents/numpyarray.py index ef9eb8a2e4..25e85ecf3c 100644 --- a/src/awkward/contents/numpyarray.py +++ b/src/awkward/contents/numpyarray.py @@ -324,6 +324,7 @@ def _is_getitem_at_placeholder(self) -> bool: @trace_function_calls def _getitem_at(self, where: IndexType): + print(" numpyarray::_getitem_at", self, where) if not self._backend.nplike.known_data and len(self._data.shape) == 1: self._touch_data(recursive=False) return TypeTracerArray._new(self._data.dtype, shape=()) @@ -332,14 +333,16 @@ def _getitem_at(self, where: IndexType): out = self._data[where] except IndexError as err: raise ak._errors.index_error(self, where, str(err)) from err - + print(" >>> out", out) if hasattr(out, "shape") and len(out.shape) != 0: + print(" >>> out has shape!", out) return NumpyArray(out, parameters=None, backend=self._backend) else: return out @trace_function_calls def _getitem_range(self, start: IndexType, stop: IndexType) -> Content: + print(" numpyarray::_getitem_range start stop", start, stop) try: out = self._data[start:stop] except IndexError as err: diff --git a/tests-cuda/test_3140_cuda_slicing.py b/tests-cuda/test_3140_cuda_slicing.py index c26f8aea3c..86ba221718 100644 --- a/tests-cuda/test_3140_cuda_slicing.py +++ b/tests-cuda/test_3140_cuda_slicing.py @@ -466,17 +466,20 @@ def cleanup_cuda(): # [{"x": [11], "y": [[None]]}], # ] +arrg_cpu = ak.Array([[1, 2, 3], [0], [4, 5]]) +arrg_gpu = ak.Array([[1, 2, 3], [0], [4, 5]], backend="cuda") + def test_simple_slice_cpu(): - arrg = ak.Array([[1, 2, 3], [0], [4, 5]]) - out = arrg[:, 0] + print("slice on CPU") + out = arrg_cpu[:, 0] expected = [1, 0, 4] result = out.tolist() cp.testing.assert_array_list_equal(result, expected, err_msg=f"Slice of [[1, 2, 3], [0], [4, 5]] should be {expected}, but got {result}") def test_simple_slice_gpu(): - arrg = ak.Array([[1, 2, 3], [0], [4, 5]], backend="cuda") - out = arrg[:, 0] + print("slice on GPU") + out = arrg_gpu[:, 0] expected = [1, 0, 4] result = out.tolist() cp.testing.assert_array_list_equal(result, expected, err_msg=f"Slice of [[1, 2, 3], [0], [4, 5]] should be {expected}, but got {result}")