numpy 4D array advanced indexing with example

Question

I am reading some deep learning code. I have problem on advanced indexing in numpy array. The code I was testing:

import numpy

x = numpy.arange(2 * 8 * 3 * 64).reshape((2, 8, 3, 64))
x.shape

p1 = numpy.arange(2)[:, None]
sd = numpy.ones(2 * 64, dtype=int).reshape((2, 64))
p4 = numpy.arange(128 // 2)[None, :]

y = x[p1, :, sd, p4]
y.shape

Why is the shape of y was (2, 64, 8)?

Here is the output of the above code:

>>> x.shape
(2, 8, 3, 64)

>>> p1
array([[0], [1]])

>>> sd
array([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
       [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])

>>> p4
array([[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15,
        16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31,
        32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47,
        48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63]])

>>> y.shape
(2, 64, 8)

I read this: https://docs.scipy.org/doc/numpy/reference/arrays.indexing.html#advanced-indexing

I think it is related to broadcasting:

x shape is (2, 8, 3, 64).

p1 is simple, it is array([[0], [1]]), just means selecting ind 0, 1 of 1st dimension. and the double array is for broadcasting.

p2 is :, which means select all 8 elements in 2nd dimension.

p3 is tricky, it contains two “lists” to pick one out of 3 elements in dimension 3, so the resulting new 3rd dimension should be 1.

p4 means it selects all 64 elements in 4th dimension.

So I think the y.shape should be (2, 8, 1, 64).

But the correct one is (2, 64, 8). Why?

Answer 1

I had the same issue when I first encountered fancy indexing in numpy. The short answer is that there is no trick to it: fancy indexing just selects elements into an output of the same shape as the index. With purely fancy indexing, your output array will be the same shape as your broadcasted index arrays (described here). The shape of the output has almost nothing to do with the shape of the input unless you throw in a regular slice index as well (described here). Your case is the latter, which adds to the confusion.

Let's take a look at your indices to see what is happening:

y = x[p1, :, sd, p4]
x.shape -> 2, 8, 3, 64
p1.shape -> 2, 1
sd.shape -> 2, 64
p4.shape -> 1, 64

The specific documentation on how to proceed is here:

Two cases of index combination need to be distinguished:

• The advanced indexes are separated by a slice, Ellipsis or newaxis. For example x[arr1, :, arr2].
• The advanced indexes are all next to each other. For example x[..., arr1, arr2, :]but not x[arr1, :, 1]since 1 is an advanced index in this regard.

In the first case, the dimensions resulting from the advanced indexing operation come first in the result array, and the subspace dimensions after that. In the second case, the dimensions from the advanced indexing operations are inserted into the result array at the same spot as they were in the initial array (the latter logic is what makes simple advanced indexing behave just like slicing).

^{Emphasis mine}

Keep in mind that in both cases described above, the dimensions of the fancy-indexed portion are the dimensions of the index arrays, not the array you are indexing.

What you should expect to see, then, is something with the broadcasted dimensions of p1, sd and p4 (2, 64), followed by the size of the second dimension of x (8). And that is indeed what you get:

>>> y.shape
(2, 64, 8)