Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[WIP] Advanced Indexing #3: Added support for multiple multidimensional Indices #8912

Closed
wants to merge 10 commits into from
35 changes: 22 additions & 13 deletions numba/core/typing/arraydecl.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,46 +33,55 @@ def get_array_index_type(ary, idx):
right_indices = []
ellipsis_met = False
advanced = False
has_integer = False

if not isinstance(idx, types.BaseTuple):
idx = [idx]

in_subspace = False
num_subspaces = 0
array_indices = []

# Walk indices
for ty in idx:
if ty is types.ellipsis:
if ellipsis_met:
raise NumbaTypeError("only one ellipsis allowed in array index "
raise NumbaTypeError("Only one ellipsis allowed in array indices "
"(got %s)" % (idx,))
ellipsis_met = True
in_subspace = False
elif isinstance(ty, types.SliceType):
pass
in_subspace = False
elif isinstance(ty, types.Integer):
# Normalize integer index
ty = types.intp if ty.signed else types.uintp
# Integer indexing removes the given dimension
ndim -= 1
has_integer = True
if not in_subspace:
num_subspaces += 1
in_subspace = True
elif (isinstance(ty, types.Array) and ty.ndim == 0
and isinstance(ty.dtype, types.Integer)):
# 0-d array used as integer index
ndim -= 1
has_integer = True
if not in_subspace:
num_subspaces += 1
in_subspace = True
elif (isinstance(ty, types.Array)
and ty.ndim == 1
and isinstance(ty.dtype, (types.Integer, types.Boolean))):
if advanced or has_integer:
# We don't support the complicated combination of
# advanced indices (and integers are considered part
# of them by Numpy).
msg = "only one advanced index supported"
raise NumbaNotImplementedError(msg)
array_indices.append(ty.ndim)
advanced = True
ndim -= 1
if not in_subspace:
num_subspaces += 1
in_subspace = True
else:
raise NumbaTypeError("unsupported array index type %s in %s"
raise NumbaTypeError("Unsupported array index type %s in %s"
% (ty, idx))
(right_indices if ellipsis_met else left_indices).append(ty)

if advanced:
ndim += max(array_indices)

# Only Numpy arrays support advanced indexing
if advanced and not isinstance(ary, types.Array):
return
Expand Down