BUG: Ensure einsum uses chunking (now that nditer doesn't) #28043
Conversation
The nditer refactor enabled the `GROWINNER` for reductions, this was not the case before (I am not sure about cases where the reduction is in an outer axis, but I don't think so). Either way, as the test says, chunking always improves the precision of large sorts if their mean is nonzero and sklearn noticed lower precision. We could only remove growinner if there is a reduction, but it seems like a 1% performance hit for the simplest (non-trivial) case where GROWINNER had done something before.
|
Thanks a lot for looking at this! I can confirm that this PR fixes scikit-learn/scikit-learn#30509. |
numpy/_core/tests/test_einsum.py
Outdated
| """ | ||
| num = 100000000 | ||
| value = 1.00000000002 | ||
| res = np.einsum("i->", np.full(num, value)) / num |
There was a problem hiding this comment.
I'd tend to use float32 and fewer elements, just to keep the test fast (now 147 ms on my machine).
There was a problem hiding this comment.
Doesn't work :). Einsum uses double for the accumulation, so for float32, the results are actually better without chunking...
But, let me use broadcasting and make the number smaller, certainly doesn't need this much.
| @@ -1035,7 +1035,6 @@ PyArray_EinsteinSum(char *subscripts, npy_intp nop, | |||
| iter_flags = NPY_ITER_EXTERNAL_LOOP| | |||
| NPY_ITER_BUFFERED| | |||
| NPY_ITER_DELAY_BUFALLOC| | |||
| NPY_ITER_GROWINNER| | |||
There was a problem hiding this comment.
How about adding a comment here? Partially as a reminder that a pair-wise sum might be a good idea...
There was a problem hiding this comment.
Yeah, didn't think of the "let's do pairwise" spin, so wasn't sure.
|
Great, thanks a lot for the investigation and the fix! Maybe one day I will manage to understand these tricky floating point things, maybe one day 😅. |
Conceptually, (in decimals) it comes down to: But if you chunk it up into batches of 10 |
The nditer refactor enabled the
GROWINNERfor reductions, this was not the case before (I am not sure about cases where the reduction is in an outer axis, but I don't think so).Either way, as the test says, chunking always improves the precision of large sums if their mean is nonzero and sklearn noticed lower precision.
We could only remove growinner if there is a reduction, but it seems like a 1% performance hit for the simplest (non-trivial) case where GROWINNER had done something before.
I suspect this is a pragmatic thing (little downside, better precision for large arrays). An even better thing would be pairwise summation, within einsum I suppose.
I.e. the reason for the chunking used to be that the iterator didn't manage to optimize this (an explicit code comment). But now, it became a choice to help the sum precision a bit...
See also scikit-learn/scikit-learn#30509, ping @lesteve