RFC: add support for `nan*` reductions

[betatim]

Is there advice around handling NaNs and how to translate Numpy code to using the Array API?

In particular I have code like np.nanmin(X, axis=0) that I would like to rewrite so that it works with Numpy, Torch, etc arrays.

To me the "obvious" translation seems to be xp.min(X[~xp.isnan(X)], axis=0) but this doesn't work if X has a shape like (7500, 4) (here xp is the namespace of the array X). Another option I looked for is a where= argument to min(), but that doesn't exist unfortunately.

Does anyone have advice on this topic or knows if there is work planned on this?

[rgommers]

Good question, thanks @betatim. I think we want a general pattern here that can apply to any reduction. That's a bit of an unsolved problem even in NumPy, where we rejected more niche nan* feature requests like nanptp and nancov.

This discussion seems relevant: pytorch/pytorch#38349 (comment). SciPy's nan_policy is now in better shape, but I'm still not sure on what the best approach in general is. Options include:

• a skip_nan keyword
• a context manager or other generic method to do achieve the same effect
• a compat library approach with a fast path for libraries known to offer specific functions like nanmin and a generic fallback implementation that iterates over the specified axis=.

A generic implementation is pretty fiddly

def nanreduction(func, x, axis=None):
    if axis is None:
        return func(x[~np.isnan(x)])
    else:
        # normalize axis, then:
        out_shape = list(x.shape).remove(x.shape[axis])
        out = np.empty_like(x, shape=out_shape)
        for i in range(axis):
            # call `func` in a loop here. this is pretty annoying to get right
            ...

NumPy's machinery is pure Python so can be looked at for guidance. It does also raise questions like how to deal with all-nan slices (e.g., https://github.com/numpy/numpy/blob/main/numpy/lib/nanfunctions.py#L355-L360); those warnings tend to not be useful.

[ogrisel]

how to deal with all-nan slices

I think the reduction should allow to pass an explicit scalar value (e.g. typically 0.) as a kwarg and raise an error if all-nan slices occur while this kwarg was left to None.

[rgommers]

I think the reduction should allow to pass an explicit scalar value (e.g. typically 0.) as a kwarg and raise an error if all-nan slices occur while this kwarg was left to None.

I'm not sure whether the better default is to return nan or raise an exception, but it should be one or the other I think. These warnings are quite annoying:

>>> np.nanmin([np.nan])
RuntimeWarning: All-NaN axis encountered
...
nan

[rgommers]

I think I had a better idea:

def nanreduction(func : Callable) -> Callable:
    # Takes any reduction in the standard as input, and returns a `nan<reduction>`
    # function that does the expected thing. This is fully functional and has a
    # small API surface, and as a bonus would be more general than the 
    # 14 `nan*` functions that numpy has (numpy has rejected requests for other
    # nan-functions because of the impact on API size).

# Usage equivalent to `nanmax(x)`:
nanreduction(xp.max)(x)

The implementation would be quite straightforward, it could be a simple dict lookup that maps max to nanmax, and so on.

@betatim @ogrisel WDYT?

[ogrisel]

Sounds good to me but ideally with an extra parameter to control the scalar value used in case of all nans. The default behavior could be to raise an exception.

[rgommers]

but ideally with an extra parameter to control the scalar value used in case of all nans

Ugh, this seems like a pretty obvious bug in NumPy:

>>> np.max([])
ValueError: zero-size array to reduction operation maximum which has no identity
>>> np.max([], initial=3)
3.0

>>> np.nanmax([])
ValueError: zero-size array to reduction operation maximum which has no identity
>>> np.nanmax([], initial=3)
<ipython-input-13-63ff44c70a84>:1: RuntimeWarning: All-NaN axis encountered
  np.nanmax([], initial=3)
nan

Other than that, I think your request is basically to support the initial keyword in all reductions, right?

[ogrisel]

Other than that, I think your request is basically to support the initial keyword in all reductions, right?

Indeed, I did not know about initial. I think the warning could be omitted.

[ogrisel]

@betatim you want want to express your opinion on the proposed API above :)

[betatim]

Sounds reasonable to me. I'm not sure I understand the comment:

The implementation would be quite straightforward, it could be a simple dict lookup that maps max to nanmax, and so on.

The way I understand the proposal is that the Array API provides nanreduction and that it is left to the user to use that together with xp.max to make themselves a nanmax.

[rgommers]

@betatim what I meant was that the implementation in numpy could look something like:

_nanreductions_mapping = dict(
    np.max: np.nanmax,
    np.min: np.nanmin,
    np.mean: np.nanmean,
     # and so on for all existing nan* funcs
)
def nanreduction(func : Callable) -> Callable:
    try:
        return _nanreductions_mapping[func]
    except KeyError:
        # generic implementation, create new function that:
        # 1. creates array with correct output shape
        # 2. per axis, fills output array with func(x[~isnan(x)])

[asmeurer]

The functional approach feels a little odd to me. Functions that return other functions aren't typically used in Python, except as decorators. Is the idea that this could in principle be applied to a user defined reduction?

I'm honestly not seeing why this is better than just adding a keyword to the various reductions.

For empty/all nan reductions, I think it should behave the same as the normal reduction. For min and max, we specify that it's implementation-defined, https://data-apis.org/array-api/latest/API_specification/generated/array_api.max.html, but for mean and std we say it must be nan https://data-apis.org/array-api/draft/API_specification/generated/array_api.mean.html https://data-apis.org/array-api/draft/API_specification/generated/array_api.std.html.

[rgommers]

The functional approach feels a little odd to me. Functions that return other functions aren't typically used in Python, except as decorators. Is the idea that this could in principle be applied to a user defined reduction?

Yes, and also to reductions in numpy et al. that do not have a nanxxx variant.

I'm honestly not seeing why this is better than just adding a keyword to the various reductions.

The response I expect to that from numpy folks at least (and it'd be my own too) is "why a keyword, we already covered this with separate functions". If it's a more general API instead, the case is that it can apply to all reductions (including those outside of numpy), and there's a reason to introduce it beyond array API standard support.