索引时如何解包元组？

Question

I have a very high dimensional tensor, say A with shape 5 X 10 X 100 X 200 X 50. I have a some numpy expression that returns a tuple T, containing indices of elements that I want to extract from A.

I'm trying this:

A[*T]

It says:

invalid syntax, you cannot use starred expressions here.

How can I do it? PS: The long solution is: A[T[0], T[1], T[2], T[3], T[4]]

EDIT: I just found that there is no need to do that as it is being done automatically. Example:

a= np.random.rand(3,3)
a[np.triu_indices(3)]

The expression np.triu_indices(3) is being unpacked automatically when passed to a as index. However , going back to my question it is not happening. To be concrete, here's an example:

a = np.random.rand(100, 50, 14, 14)
a[:, :, np.triu_indices(14)].shape

Supposedly, the last bit np.triu_indices(14) should act on last two axes, as in the previous example, but it is not happening, and the shape resulting is weird. Why isn't being unpacked? and how to do that?

Answer 1

The problem with:

a[:, :, np.triu_indices(14)]

is that you are using as argument for [] a tuple of mixed types slice and tuple ( tuple(slice, slice, tuple(np.ndarray, np.ndarray)) ) and not a single tuple (eventually with advanced indexing), eg tuple(slice, slice, np.ndarray, np.ndarray) . This is causing your troubles. I would not go into the details of what is happening in your case.

Changing that line to:

a[(slice(None),) * 2 + np.triu_indices(14)]

will fix your issues:

a[(slice(None),) * 2 + np.triu_indices(14)].shape
# (100, 50, 105)

---

Note that there are a couple of ways for rewriting:

(slice(None),) * 2 + np.triu_indices(14)

another way may be:

(slice(None), slice(None), *np.triu_indices(14))

---

Also, if you want to use the ... syntax, you need to know that ... is syntactic sugar for Ellipsis , so that:

(Ellipsis,) + np.triu_indices(14)

or:

(Ellipsis, *np.triu_indices(14))

would work correctly:

a[(Ellipsis,) + np.triu_indices(14)].shape                                                   
# (100, 50, 105)
a[(Ellipsis, *np.triu_indices(14))].shape                                                   
# (100, 50, 105)