What was Jax's idea behind implementing the C standard Annex G recoveries

[ikrommyd]

Hi y'all,

I'm interested in what was Jax's logic behind implementing the Annex G recoveries of the C standard for complex division.
See numpy/numpy#30806 and the linked issues for a lot more context.
In that NumPy PR, I had the intention of adding the Annex G recoveries to recover NaNs in complex number multiplication and division given that CPython somewhat recently added them too (issues linked there as well).
However in NumPy, we decided to eventually not do it now, especially for complex multiplication as it was adding a lot more overhead in the nominal case. For division the overhead I was seeing was like at the few % level.
However, after closing the PR, I noticed that Jax does indeed implement the Annex G recoveries for complex division and not for multiplication on cpu and gpu. Cupy does not implement them at all.
So my question I guess is, what was Jax's logic behind that? Was the argument that a small performance hit is fine for correctness of NaNs?
I'm trying to find arguments whether we should add them in NumPy or not for division only.

Cheers!

[hawkinsp]

The relevant code lives inside XLA. My strong suspicion is that this was done to match the behavior of the C++ standard library for these values, rather than with any great consideration for a performance/correctness tradeoff.

(And in general, absent such reasoning about a tradeoff, I'd argue that "correctness" is a great place to start.)

[hawkinsp]

I do however suspect that the cost of the annex G recovery is relatively small for division but relatively expensive for multiplication. There's just much less to hide behind in multiplication.

[ikrommyd]

Yeah I don't remember what the standard library does actually. Very likely they implement it for multiplication too.
I had done some pyperf benchmarks in numpy/numpy#30806 and what I saw is that for multiplication it was annoying even when there were no NaNs. Cause multiplication is SIMD accelerated in numpy, the extra isnan check after the SIMD pass was a killer. For division it was not that important since division is expensive and not SIMD accelerated as well. I wanted to revisit it to perhaps do it only for division in numpy and maybe then cupy can be convinced to do it too. For multiplication I'd argue it's not worth it for scientific computing due to the performance hit.