💪 handle NaNs

[jvdd]

We want to properly handle NaNs.

What are NaNs?

NaNs are only present in float datatypes - and thus do not exist in (unsigend) integer datatypes.
For simplicity I will illustrate things for 16-bit numbers.

Some background ⬇️

Int16 representation

https://steffanynaranjo.medium.com/how-integers-are-stored-in-memory-using-twos-complement-8b4237f7426d

Float16 representation

!! We are dealing with fp16 here, NOT bfloat16

numpy.float16: 16-bit-precision floating-point number type: sign bit, 5 bits exponent, 10 bits mantissa.

So how do NaNs look like?

nans occur when exponent is all 1s and mantissa is not all 0s. (sign does not matter)

How do infinities look like?

+inf occurs when sign is 0, exp is all 1s, and mantissa is all 0s.

-inf occurs when sign is 1, exp is all 1s, and mantissa is all 0s.

How does the current implementation cope with NaNs?

Necessary background info: every comparison (gt, lt, eq) with NaNs results in false ⬇️

let v = f32::NAN;
assert_eq!(v > 5.0, false);
assert_eq!(v < 5.0, false);
assert_eq!(v == 5.0, false);

=> Current implementation deals as follows with NaNs in the data:

• if there are NO NaNs in the first SIMD lane: current implementation will ignore the NaNs (as comparisons with NaNs result in false and thus no NaNs will be added in the accumulating SIMD vector)
  • this corresponds to the np.nanargmax / np.nanargmin behavior
• if there are ONE or MULTIPLE NaNs in first SIMD lane: current implementation will never update the NaNs in the accumulating SIMD vector (as comparisons will result in false and thus the NaNs will never be updated).

TODOs

• fix this issue with NaN values in the first SIMD vector -> if we can create/ensure a "NaN-free" initial SIMD vec, than no NaNs will be added in the inner SIMD loop => we have an implementation of np.nanargmin / np.nanargmax.

How should we handle NaNs

Ideally we support, just like numpy, two variants of the argmim / argmax algorithm:

1. ignoring NaNs: as the current version does (if we fix the TODO above)
   -> corresponds to np.nanargmin / np.nanargmax
2. returning first NaN index once there is one present in the data
   -> corresponds to np.argmin / np.argmax

we can serve both functionalities by adding a new function to the ArgMinMax trait and let it default for non-float datatypes to the current implementation.

[jvdd]

What I tried in PR #18

Transform floats to integer dtypes & perform comparisons in that data space

How: use same projection (bitwise transformation & transmutation) as we did for f16, see #1

The projection

ord_transform(v: i16) = ((v >> 15) & 0x7FFF) ^ v)
(to apply this on a f16, just transmute f16 to i16 first)

⬇️ illustration of how this projection (mapping float16 to int16)

• preserves ordinality
• maps +inf & -inf respectively above & below all "normal" floats
• maps NaNs to the ends of the space

Other really nice to have features:

• projection is symmetric (or bijective): performing the transformation twice places us back at the original floating value (as long as we don't change the integer value of course)

TODOs

• can we do this projection more efficiently? (e.g., less operations)
• are there other possible projections?

Some thoughts

• works fine when doing argminmax (i.e., argmin & argmax in 1 function) as we have to check both gt and lt - and thus will always find the NaN when one is present - independent of the sign bit
• when doing argmin or argmax we check only for lt or gt respectively - and thus will NOT always find the NaN when one is present
  • we could possible shift the integer projection with a wrapping_add, placing all the NaNs at the right side of the comparison - DISCLAIMER: not sure if this will work & even is possible
    • check if this is possible

[jvdd]

@varon - I just updated this Issue (added a description + explained in the comment above what I tried in PR #18)

If you can find the time, I would love to hear your feedback on this! :)

---

Some futher (optional) questions that would be interesting to discuss:

• Do you agree with adding support for "nan-omitting" & "nan-returning" argminmax?
• What is your opinion on the projection solution in PR 🚧 POC - support NaNs for SSE & AVX2 f32 #18 - see comment above? Do you think we can optimize the transformation or use something similar?
• Do you have any alternatives in mind for the "nan-returning" argminmax?