Skip to content

Commit

Permalink
fixed the logic for transposing base array when bool arrays are passe…
Browse files Browse the repository at this point in the history 
…d as indices
  • Loading branch information
@ipdemes
ipdemes committed Apr 19, 2022
1 parent 7a33d9e commit 8e6b2bb
Show file tree
Hide file tree
Showing 2 changed files with 94 additions and 26 deletions.
18 changes: 14 additions & 4 deletions cunumeric/deferred.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -465,16 +465,25 @@ def _create_indexing_array(self, key, is_set=False):
if np.isscalar(k) or isinstance(k, NumPyThunk):
if start_index == -1:
start_index = dim
transpose_indices += (dim,)
key_transpose_indices += (dim,)
transpose_needed = transpose_needed or ((dim - last_index) > 1)
if (
isinstance(k, NumPyThunk)
and k.dtype == np.bool
and k.ndim >= 2
):
for i in range(dim, dim + k.ndim):
transpose_indices += (shift + i,)
shift += k.ndim - 1
else:
transpose_indices += (dim,)
last_index = dim

if transpose_needed:
start_index = 0
post_indices = tuple(
i for i in range(store.ndim) if i not in transpose_indices
)
key_transpose_indices = transpose_indices
transpose_indices += post_indices
post_indices = tuple(
i for i in range(len(key)) if i not in key_transpose_indices
Expand All @@ -483,8 +492,8 @@ def _create_indexing_array(self, key, is_set=False):
store = store.transpose(transpose_indices)
key = tuple(key[i] for i in key_transpose_indices)

for d, k in enumerate(key):
dim = d
shift = 0
for dim, k in enumerate(key):
if np.isscalar(k):
if k < 0:
k += store.shape[dim + shift]
Expand All @@ -506,6 +515,7 @@ def _create_indexing_array(self, key, is_set=False):
# in case of the mixed indises we all nonzero
# for the bool array
k = k.nonzero()
shift += len(k) - 1
tuple_of_arrays += k
else:
tuple_of_arrays += (k,)
Expand Down
102 changes: 80 additions & 22 deletions tests/advanced_indexing.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -317,25 +317,17 @@ def test():
res_num = x_num[..., indx_bool_num]
assert np.array_equal(res, res_num)

print("IRINA DEBUG 1")
indx1_bool = np.array([True, False])
indx1_bool_num = num.array(indx1_bool)
indx2_bool = np.array([True, False, True, True])
indx2_bool_num = num.array(indx2_bool)
res = x[indx1_bool, :, indx2_bool]
print(res.shape)
print(res)
res_num = x_num[indx1_bool_num, :, indx2_bool_num]
print(res_num.shape)
print(res_num)
assert np.array_equal(res, res_num)

print("IRINA DEBUG 2")
res = x[indx1_bool, 1, indx2_bool]
# res_num = x_num[indx1_bool_num, 1, indx2_bool_num]
# print(res.shape)
# print(res_num.shape)
# assert np.array_equal(res, res_num)
res_num = x_num[indx1_bool_num, 1, indx2_bool_num]
assert np.array_equal(res, res_num)

# g: boolean array with the same shape is passed to x:
indx = x % 2
Expand All @@ -351,15 +343,13 @@ def test():
z[indx] = 1
z_num[indx_num] = 1
assert np.array_equal(z, z_num)
print("IRINA DEBUG 3")

indx_bool = np.array([True, False, True])
indx_bool_num = num.array(indx_bool)
z[:, indx_bool] = 5
z_num[:, indx_bool_num] = 5
assert np.array_equal(z, z_num)

print("IRINA DEBUG 4")
# i: two bool array of the same shape are passed:
x = mk_seq_array(
np,
Expand Down Expand Up @@ -388,20 +378,66 @@ def test():
)
indx_num = num.array(indx)
res = x[indx, indx]
print("IRINA DEBUG res = ", res.shape)
# res_num = x_num[indx_num, indx_num]
# assert np.array_equal(res, res_num)
res_num = x_num[indx_num, indx_num]
assert np.array_equal(res, res_num)

x = mk_seq_array(
np,
(
3,
4,
5,
3,
4,
),
)
x_num = mk_seq_array(
num,
(
3,
4,
5,
3,
4,
),
)
res = x[indx, 1, indx]
res_num = x_num[indx_num, 1, indx_num]
assert np.array_equal(res, res_num)

res = x[indx, :, indx]
res_num = x_num[indx_num, :, indx_num]
assert np.array_equal(res, res_num)

# j: 2 bool arrays should be broadcasted:
# res = x[idx, [True,False,False]]
# res_num = x_num[idx_num, [True,False,False]]
x = mk_seq_array(
np,
(
3,
4,
3,
4,
),
)
x_num = mk_seq_array(
num,
(
3,
4,
3,
4,
),
)
res = x[indx, [True, False, False]]
res_num = x_num[indx_num, [True, False, False]]
assert np.array_equal(res, res_num)

# 2d bool array not at the first index:
indx = np.full((4, 3), True)
indx_num = num.array(indx)
res = x[:, indx]
# res_num = x_num[:, indx]
# assert np.array_equal(res, res_num)
res_num = x_num[:, indx]
assert np.array_equal(res, res_num)

# 3: testing mixed type of the arguments passed:

Expand All @@ -425,8 +461,8 @@ def test():
),
)
res = x[[1, 1], [False, True, False]]
# res_num = x_num[[1,1], [False, True,False]]
# assert np.array_equal(res, res_num)
res_num = x_num[[1, 1], [False, True, False]]
assert np.array_equal(res, res_num)

res = x[[1, 1], :, [False, True, False, True]]
res_num = x_num[[1, 1], :, [False, True, False, True]]
Expand Down Expand Up @@ -465,7 +501,7 @@ def test():
res_num = x_num[:, [0, 1], :, 1:]
assert np.array_equal(res, res_num)

# c: transformed base:
# c: transformed base or index or rhs:
z = x[:, 1:]
z_num = x_num[:, 1:]
indx = np.array([1, 1])
Expand All @@ -474,6 +510,22 @@ def test():
res_num = z_num[indx_num]
assert np.array_equal(res, res_num)

indx = np.array([1, 1, 0])
indx_num = num.array(indx)
indx = indx[1:]
indx_num = indx_num[1:]
res = z[1, indx]
res_num = z_num[1, indx_num]
assert np.array_equal(res, res_num)

b = np.ones((2, 3, 6, 5))
b_num = num.array(b)
b = b.transpose((0, 1, 3, 2))
b_num = b_num.transpose((0, 1, 3, 2))
z[indx] = b
z_num[indx_num] = b_num
assert np.array_equal(z, z_num)

# d: shape mismatch case:
x = np.array(
[
Expand Down Expand Up @@ -536,6 +588,12 @@ def test():
x_num[ind_num, ind_num] = 5
assert np.array_equal(x, x_num)

b = np.array([1, 2, 3], dtype=np.int16)
b_num = num.array(b)
x[ind, ind] = b
x_num[ind_num, ind_num] = b_num
assert np.array_equal(x, x_num)

# we do less than LEGATE_MAX_DIM becasue the dimension will be increased by
# 1 when passig 2d index array
for ndim in range(2, LEGATE_MAX_DIM):
Expand Down

0 comments on commit 8e6b2bb

Please sign in to comment.