Broadcasting pointwise indexing operator #2591
Description
I would like to suggest to add a broadcasting pointwise indexing operator to Julia. This a very expressive form of indexing, and quite handy, as remarked in #2247.
Let us call the new operation pointref(A, inds...) = A.[inds...], where A and the elements of inds are arrays. Then we can compute
X = pointref(A, inds...) = A.[inds...]
as follows:
-
Broadcast all elements of
indsto the same shape: Extrude each singleton
axis in eachindarray to the common non-singleton size of that
axis among them. It is an error if the indices cannot be broadcast together. -
The shape of
Xis the common broadcast shape of the indices. -
Each entry of
Xis calculated by indexingAwith the elements ofindsat the corresponding position. Suppose thatindshave been broadcast to the same shape. Then
for each element of the resultX[ks...],X[ks...] = A[[ind[ks...] for ind in inds]...]
Example:
A = [1 3
2 4]
P = [2 1 1]
Q = [1 2 1
2 1 2]
X = pointref(A, P, Q) = A.[P, Q]
Since size(P) = (1,3) and size(Q) = (2,3), their broadcast shape is (2,3). The broadcast index arrays are
Pb = [2 1 1
2 1 1]
Qb = [1 2 1
2 1 2]
where P has been extruded along the first axis. Then
X = [A[2,1] A[1,2] A[1,1]
A[2,2] A[1,1] A[1,2]]
= [2 3 1
4 1 3]
Relationship to current indexing
Current Julia indexing allows
A[inds...]
where the elements of inds are vectors and/or scalars. This axiswise indexing can be expressed using pointwise indexing as
A[inds...] = A.[[swapdims(ind,1,k) for (k,ind) in enumerate(inds)]...]
where swapdims(A,j,k) swaps axes j and k of an array. We see that pointwise and axiswise indexing coincide when there is only one index, since swapdims(A,1,1) is a noop.
Julia also supports linear indexing
A[ind] = A[:][ind]
Analogously, one could consider
A.[ind] = A[:].[ind]
i.e. A.[ind] looks does a linear-index lookup in A for each element in the index array ind.
This seems to be the same behavior as suggested in #2247 for A[ind].
Wrap-up
- I think that pointwise indexing is useful enough to deserve its own operator.
- It's also different enough from regular indexing that it seems impossible to overload the two in the same operator, except for the single-index case
A[B], where the two can be said to coincide. (As described in RFC: Linear indexing with multidimensional index. #2247 (comment), numpy does this also in the multi-index case, which becomes messy.) - It also happens to be a natural way to express the kind of indexing that a gpu kernel is capable of, into an array stored in gpu memory.