What is the purpose of numpy.where returning a tuple?

Question

When I run this code:

import numpy as np a = np.array([1, 2, 3, 4, 5, 6]) print(np.where(a > 2)) 

it would be natural to get an array of indices where a > 2, i.e. [2, 3, 4, 5], but instead we get:

(array([2, 3, 4, 5], dtype=int64),) 

i.e. a tuple with empty second member.

Then, to get the the "natural" answer of numpy.where, we have to do:

np.where(a > 2)[0] 

What's the point in this tuple? In which situation is it useful?

Note: I'm speaking here only about the use case numpy.where(cond) and not numpy.where(cond, x, y) that also exists (see documentation).

Answer 1

numpy.where returns a tuple because each element of the tuple refers to a dimension.

Consider this example in 2 dimensions:

a = np.array([[1, 2, 3, 4, 5, 6],               [-2, 1, 2, 3, 4, 5]])  print(np.where(a > 2))  (array([0, 0, 0, 0, 1, 1, 1], dtype=int64),  array([2, 3, 4, 5, 3, 4, 5], dtype=int64)) 

As you can see, the first element of the tuple refers to the first dimension of relevant elements; the second element refers to the second dimension.

This is a convention numpy often uses. You will see it also when you ask for the shape of an array, i.e. the shape of a 1-dimensional array will return a tuple with 1 element:

a = np.array([[1, 2, 3, 4, 5, 6],               [-2, 1, 2, 3, 4, 5]])  print(a.shape, a.ndim)  # (2, 6) 2  b = np.array([1, 2, 3, 4, 5, 6])  print(b.shape, b.ndim)  # (6,) 1