From 018fd212f40d4ef8d4ec81d9e31940c069a08c5c Mon Sep 17 00:00:00 2001 From: Edwin Date: Sun, 8 Jun 2025 11:38:49 -0700 Subject: [PATCH 1/2] Fixed indexing with Arrays --- arrayfire/array_object.py | 78 +++++++++++++++++++++++++++++++++------ 1 file changed, 66 insertions(+), 12 deletions(-) diff --git a/arrayfire/array_object.py b/arrayfire/array_object.py index d44ee29..c79458c 100755 --- a/arrayfire/array_object.py +++ b/arrayfire/array_object.py @@ -746,7 +746,7 @@ def __getitem__(self, key: IndexKey, /) -> Array: ---------- self : Array Array instance. - key : int | slice | tuple[int | slice, ...] | Array + key : int | slice | tuple[int | slice | Array, ...] | Array Index key. Returns @@ -754,19 +754,32 @@ def __getitem__(self, key: IndexKey, /) -> Array: out : Array An array containing the accessed value(s). The returned array must have the same data type as self. """ - # TODO - # API Specification - key: Union[int, slice, ellipsis, tuple[Union[int, slice, ellipsis], ...], array]. - # consider using af.span to replace ellipsis during refactoring out = Array() ndims = self.ndim - if isinstance(key, Array) and key == afbool.c_api_value: + indexing = key + + if isinstance(key, int | float | slice): # when indexing with one dimension, treat it as indexing a flat array ndims = 1 - if wrapper.count_all(key.arr) == 0: # HACK was count() method before - return out - - # HACK known issue - out._arr = wrapper.index_gen(self._arr, ndims, wrapper.get_indices(key)) # type: ignore[arg-type] + elif isinstance(key, Array): # when indexing with one array, treat it as indexing a flat array + ndims = 1 + if key.is_bool: + indexing = wrapper.where(key.arr) + else: + indexing = key.arr + elif isinstance(key, tuple): + key_list = [] + for elem in key: + if isinstance(elem, Array): + if elem.is_bool: + key_list.append(wrapper.where(elem.arr)) + else: + key_list.append(elem.arr) + else: + key_list.append(elem) + indexing = tuple(key_list) + + out._arr = wrapper.index_gen(self._arr, ndims, wrapper.get_indices(indexing)) # type: ignore[arg-type] return out def __index__(self) -> int: @@ -781,6 +794,18 @@ def __len__(self) -> int: return self.shape[0] if self.shape else 0 def __setitem__(self, key: IndexKey, value: int | float | bool | Array, /) -> None: + """ + Assigns self[key] = value + + Parameters + ---------- + self : Array + Array instance. + key : int | slice | tuple[int | slice | Array, ...] | Array + Index key. + value: int | float | complex | bool | Array + + """ ndims = self.ndim is_array_with_bool = isinstance(key, Array) and type(key) is afbool @@ -803,7 +828,29 @@ def __setitem__(self, key: IndexKey, value: int | float | bool | Array, /) -> No other_arr = value.arr del_other = False - indices = wrapper.get_indices(key) # type: ignore[arg-type] # FIXME + indexing = key + if isinstance(key, int | float | slice): # when indexing with one dimension, treat it as indexing a flat array + ndims = 1 + elif isinstance(key, Array): # when indexing with one array, treat it as indexing a flat array + ndims = 1 + if key.is_bool: + indexing = wrapper.where(key.arr) + else: + indexing = key.arr + elif isinstance(key, tuple): + key_list = [] + for elem in key: + if isinstance(elem, Array): + if elem.is_bool: + key_list.append(wrapper.where(elem.arr)) + else: + key_list.append(elem.arr) + else: + key_list.append(elem) + indexing = tuple(key_list) + + indices = wrapper.get_indices(indexing) + out = wrapper.assign_gen(self._arr, other_arr, ndims, indices) wrapper.release_array(self._arr) @@ -1144,7 +1191,10 @@ def _get_processed_index(key: IndexKey, shape: tuple[int, ...]) -> tuple[int, .. if isinstance(key, tuple): return tuple(_index_to_afindex(key[i], shape[i]) for i in range(len(key))) - return (_index_to_afindex(key, shape[0]),) + shape[1:] + size = 1 + for dim_size in shape: + size *= dim_size + return (_index_to_afindex(key, size),) def _index_to_afindex(key: int | float | complex | bool | slice | wrapper.ParallelRange | Array, axis: int) -> int: @@ -1168,6 +1218,10 @@ def _index_to_afindex(key: int | float | complex | bool | slice | wrapper.Parall def _slice_to_length(key: slice, axis: int) -> int: + start = key.start + stop = key.stop + step = key.step + if key.start is None: start = 0 elif key.start < 0: From 0b3f859a3703a31f77ab55d68bee0073abe6da0c Mon Sep 17 00:00:00 2001 From: Edwin Solis Date: Mon, 16 Jun 2025 23:26:59 -0700 Subject: [PATCH 2/2] Fixed memory leak when indexing with bools --- arrayfire/array_object.py | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/arrayfire/array_object.py b/arrayfire/array_object.py index c79458c..ea8c5cd 100755 --- a/arrayfire/array_object.py +++ b/arrayfire/array_object.py @@ -780,6 +780,14 @@ def __getitem__(self, key: IndexKey, /) -> Array: indexing = tuple(key_list) out._arr = wrapper.index_gen(self._arr, ndims, wrapper.get_indices(indexing)) # type: ignore[arg-type] + + if isinstance(key, Array) and key.is_bool: + wrapper.release_array(indexing) + elif isinstance(key, tuple): + for i in range(len(key)): + if isinstance(key[i], Array) and key[i].is_bool: + wrapper.release_array(indexing[i]) + return out def __index__(self) -> int: @@ -807,7 +815,6 @@ def __setitem__(self, key: IndexKey, value: int | float | bool | Array, /) -> No """ ndims = self.ndim - is_array_with_bool = isinstance(key, Array) and type(key) is afbool if is_array_with_bool: @@ -842,20 +849,27 @@ def __setitem__(self, key: IndexKey, value: int | float | bool | Array, /) -> No for elem in key: if isinstance(elem, Array): if elem.is_bool: - key_list.append(wrapper.where(elem.arr)) + locs = wrapper.where(elem.arr) + key_list.append(locs) else: key_list.append(elem.arr) else: key_list.append(elem) indexing = tuple(key_list) - indices = wrapper.get_indices(indexing) + out = wrapper.assign_gen(self._arr, other_arr, ndims, wrapper.get_indices(indexing)) - out = wrapper.assign_gen(self._arr, other_arr, ndims, indices) + if isinstance(key, Array) and key.is_bool: + wrapper.release_array(indexing) + elif isinstance(key, tuple): + for i in range(len(key)): + if isinstance(key[i], Array) and key[i].is_bool: + wrapper.release_array(indexing[i]) wrapper.release_array(self._arr) if del_other: wrapper.release_array(other_arr) + self._arr = out def __str__(self) -> str: