@@ -780,6 +780,14 @@ def __getitem__(self, key: IndexKey, /) -> Array:
780
780
indexing = tuple (key_list )
781
781
782
782
out ._arr = wrapper .index_gen (self ._arr , ndims , wrapper .get_indices (indexing )) # type: ignore[arg-type]
783
+
784
+ if isinstance (key , Array ) and key .is_bool :
785
+ wrapper .release_array (indexing )
786
+ elif isinstance (key , tuple ):
787
+ for i in range (len (key )):
788
+ if isinstance (key [i ], Array ) and key [i ].is_bool :
789
+ wrapper .release_array (indexing [i ])
790
+
783
791
return out
784
792
785
793
def __index__ (self ) -> int :
@@ -807,7 +815,6 @@ def __setitem__(self, key: IndexKey, value: int | float | bool | Array, /) -> No
807
815
808
816
"""
809
817
ndims = self .ndim
810
-
811
818
is_array_with_bool = isinstance (key , Array ) and type (key ) is afbool
812
819
813
820
if is_array_with_bool :
@@ -842,20 +849,27 @@ def __setitem__(self, key: IndexKey, value: int | float | bool | Array, /) -> No
842
849
for elem in key :
843
850
if isinstance (elem , Array ):
844
851
if elem .is_bool :
845
- key_list .append (wrapper .where (elem .arr ))
852
+ locs = wrapper .where (elem .arr )
853
+ key_list .append (locs )
846
854
else :
847
855
key_list .append (elem .arr )
848
856
else :
849
857
key_list .append (elem )
850
858
indexing = tuple (key_list )
851
859
852
- indices = wrapper .get_indices (indexing )
860
+ out = wrapper .assign_gen ( self . _arr , other_arr , ndims , wrapper . get_indices (indexing ) )
853
861
854
- out = wrapper .assign_gen (self ._arr , other_arr , ndims , indices )
862
+ if isinstance (key , Array ) and key .is_bool :
863
+ wrapper .release_array (indexing )
864
+ elif isinstance (key , tuple ):
865
+ for i in range (len (key )):
866
+ if isinstance (key [i ], Array ) and key [i ].is_bool :
867
+ wrapper .release_array (indexing [i ])
855
868
856
869
wrapper .release_array (self ._arr )
857
870
if del_other :
858
871
wrapper .release_array (other_arr )
872
+
859
873
self ._arr = out
860
874
861
875
def __str__ (self ) -> str :
0 commit comments