ENH: Add a where argument to np.mean

This would match and be implemented via the where argument in np.sum.

Most of the work here comes from implementing np.core._methods._count_reduce_items correctly to handle the where argument.

Perhaps worth fixing #15817 first.

[sgasse]

@eric-wieser sure, I just did not want to waste CI resources for a PR that I know would fail the checks.

You are right that speed is a limitation. In this commit, I catch the case that where is just a boolean and do not initialize the logical mask in that case. This gives a good speedup for cases where the parameter is not set explicitly. I have also added a few test cases and made sure to pass all other tests without changing them.

I think optimization potential lies in the ‘in between’ cases: When where is a broadcastable array that is much smaller than the array of which to calculate the mean. Maybe I can find an abstract way of calculating rcount for those cases without a mask.

Other than that, do you think there is a much faster way of calculating the mean in case a dense mask is given as where?

[mtsokol]

Hi! I’ve been porting this enhancement to jax and just wanted to ask if such error with where keyword is expected: https://github.com/numpy/numpy/issues/18552 Thank you for any advice!