Specify NaN behaviour in `unique()`

[honno]

Say we have an array with multiple NaN values:

>>> x = xp.full(5, xp.nan)
>>> x
Array([nan, nan, nan, nan, nan])

Should NaNs be treated as unique to one another?

>>> xp.unique(x)
Array([nan, nan, nan, nan, nan])

Or should they be treated as the same?

>>> xp.unique(x)
Array([nan])

I have no dog in the race but you might want to refer to some discussion bottom of numpy/numpy#19301 and in the NumPy mailing list that relates to a recent accidental "regression" in how np.unique() deals with NaNs.

In either case, just specificity would prevent creating wrapper methods to standardise this behaviour. For example I created this (admittedly scrappy) utility method for HypothesisWorks/hypothesis#3065 to work for both scenarios, which @asmeurer noted was not ideal and probably a failure of the Array API spec.

def count_unique(array):
    n_unique = 0
    nan_index = xp.isnan(array)
    for isnan, count in zip(*xp.unique(nan_index, return_counts=True)):
        if isnan:
            n_unique += count
            break
    filtered_array = array[~nan_index]
    unique_array = xp.unique(filtered_array)
    n_unique += unique_array.size
    return n_unique

And if deterministic NaN behaviour cannot be part of the API, a note should be added to say this behaviour is out of scope.

[pmeier]

IMO this depends on the definition of "unique". If it means "value equality", it should return all occurrences of NaN, since assert float("nan") != float("nan"). If it means "identity", it should return only one NaN, since assert float("nan") is float("nan"). (See #249 (comment))

If I'm not mistaken, in this context "unique" means value equality, since we are dealing with values and not objects (is np.unique defined for the object dtype?). To satisfy both requirements, we could add an equal_nan flag that, if True, treats NaN's as equal. This is similar to what most testing functions do, e.g. np.testing.assert_allclose.

[rgommers]

Are there clear use cases for returning multiple nan's? As @thomasjpfan said in numpy/numpy#19301, a single nan is what is expected for ML applications. And in other contexts I'd expect that to be true as well. I can't remember any case where I had a need for multiple distinct nan's.

[asmeurer]

How do pytorch and tensorflow (and any others) handle this? Do they also have discussions about this behavior?

[pmeier]

For PyTorch I'm not aware of a discussion about this. As for the behavior, it seems to use the value equality approach I described in my comment above:

>>> t = torch.tensor([1.0, float("nan"), 1.0, float("nan")])
>>> torch.unique(t)
tensor([nan, nan, 1.])

[Zac-HD]

assert float("nan") is float("nan")

>>> float('nan') is float('nan')
False

math.nan is math.nan, sure, but otherwise identity of separately-constructed objects is implementation-defined and probably False. The only reliable exceptions are singletons (None, False, True, etc.) and some implementation-defined details like CPython's interning for small integers and strings.

(personally I favor the "return one nan" option, but we'd better document whatever we decide)

[pmeier]

@Zac-HD my bad, thanks for the heads-up. I've fixed my comment above.

[rgommers]

TensorFlow and CuPy also return multiple nan's. CuPy will almost certainly change that to match NumPy. For TensorFlow I also could not find an issue. It seems likely that in all libraries the multiple-nan return is just a natural result of implementing unique via comparisons without any special handling, not because it's useful behaviour.

[leofang]

It'd be very annoying to implement on GPUs if we were to return a unique nan. In any case unique relies on sort or argsort, and it creates a thread divergence that hurts performance just to check special cases like this. We were (=I was) bothered enough to follow NumPy's way to handle complex nan ordering, and would prefer to not add any more special cases. (You can imagine it'd only hurt performance more by an additional call to isnan just to mark where the nan starts in a sorted array, count the occurrence so as to allocate output memory, etc)

[rgommers]

Hmm, that's a good point. Although it raises the question what users do then, they usually need to handle the nan's anyway by skipping or combining them.

Your argument means adding an equal_nan also doesn't help, then you'd still have to implement it.

[honno]

Could equal_nan be a kwarg which must support None (i.e. whatever xp.unique() would do by default) but can omit either/both True and False? The spec doesn't allow omission of any specific arguments per-say but it does allow omission of boolean indexing (technically a specified argument for get/set item), so maybe it's not too "magical" of a compromise.

It could even be that a library which knows how it will return NaNs in xp.unique() can allow the respective equal_nan bool arg to just work, and raise a helpful error for the opposite arg.

[jakirkham]

As an aside, we discovered this behavior on the call today.

In [1]: import numpy as np

In [2]: {float('nan'), float('nan')}
Out[2]: {nan, nan}

In [3]: {np.nan, np.nan}
Out[3]: {nan}

[jakirkham]

From a suggestion in the call, what if a unique implementation filtered out nans first? Something like this?

a_notnan = a[~np.isnan(a)]
r = np.unique(a_notnan)

[pmeier]

@jakirkham

As an aside, we discovered this behavior on the call today.

In [1]: import numpy as np

In [2]: {float('nan'), float('nan')}
Out[2]: {nan, nan}

In [3]: {np.nan, np.nan}
Out[3]: {nan}

That is the same trap I fell into: Since np.nan is only generated once, it is by design the identical object and thus will be filtered by the set:

In[1]: nan = float("nan")

In[2]: {nan, nan}
Out[2]: {nan}

[jakirkham]

Yeah was thinking that might be the case. Thanks for sharing Philip 🙂

[honno]

From a suggestion in the call, what if a unique implementation filtered out nans first? Something like this?

a_notnan = a[~np.isnan(a)]
r = np.unique(a_notnan)

A limitation with this example specifically is that boolean indexing might not be supported by an Array API library (it's optional behaviour in the spec, see #84). You could manually zip an array a with its NaN mask xp.isnan(a) to then construct a unique array, but this could be painfully slow, and generally an awkward implementation (see Leo Fang's previous comment).