How can I do indexing in numpy over multiple samples?

Question

I have a train dataset traind with the shape (60000,28,28). I want to slice out all samples (60000) which have a mean of 0.3 or bigger. I tried this:

condition = (np.mean(traind[0:60000,:,:]) > 0.3)
subTraind = traind[condition]

but it throws me followed error:

IndexError: too many indices for array: array is 2-dimensional, but 3 were indexed

Is there a way to do a slicing with multiple index conditions?

Answer 1

np.mean will compute the mean over the entire array by default. You want to compute mean along the last 2 axes only:

condition = traind.mean(axis=(1,2)) > 0.3