Skip to content

FIX stabilize weighted confusion_matrix_at_thresholds on float32-only devices - #34827

Draft
ogrisel wants to merge 9 commits into
scikit-learn:mainfrom
ogrisel:fix-34813-cmat-weighted-int-scale
Draft

FIX stabilize weighted confusion_matrix_at_thresholds on float32-only devices#34827
ogrisel wants to merge 9 commits into
scikit-learn:mainfrom
ogrisel:fix-34813-cmat-weighted-int-scale

Conversation

@ogrisel

@ogrisel ogrisel commented Aug 27, 2026

Copy link
Copy Markdown
Member

Fixes the weighted half of #34813.

Complements #34817, which handles the unweighted case with integer cumsum.

On float32-only array API devices (e.g. array-api-strict's no_float64, torch MPS), confusion_matrix_at_thresholds accumulates weighted counts in float32 when float64 is unavailable. Float32 cumsums saturate past 2**24, which breaks ROC/PR curves and related metrics on large datasets. Weight normalization alone is not enough at that scale.

When sample_weight is set and the device only supports float32, this PR accumulates in fixed-point int64 (round(weight * 1e6)), then converts back to float32 at the output boundary.

The float64 path is unchanged, so this PR should result in a net numerical stability improvement for float32 devices. We could explore using the fixed-point code path also for devices that support float64 operations but I am not sure if this is a good idea or not.

Note to reviewers: I think we should review and merge the simpler #34817 fix first and then I can rebase this branch to review it more naturally.

ogrisel and others added 5 commits August 25, 2026 18:38
On float32-only Array API devices, float cumsum saturates past 2**24.
Weight normalization is not sufficient at that scale; use a fixed-point
integer cumulative sum instead when float64 is unavailable. Refs scikit-learn#34813.

Co-authored-by: Cursor <cursoragent@cursor.com>
Drop _max_precision_int_dtype; float32-only devices of interest
(e.g. torch MPS) provide int64.

Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Rely on _array_api_for_tests for skipping when array-api-strict or
SCIPY_ARRAY_API is unavailable.

Co-authored-by: Cursor <cursoragent@cursor.com>
@ogrisel

ogrisel commented Aug 27, 2026

Copy link
Copy Markdown
Member Author

cc @david-cortes-intel as we discussed this during a meeting.

ogrisel and others added 2 commits August 27, 2026 17:14
Add sklearn.metrics/34827.fix.rst for PR scikit-learn#34827.

Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
@ogrisel ogrisel moved this to In Progress in Array API Aug 27, 2026
@github-actions github-actions Bot removed the CUDA CI label Aug 27, 2026
Comment thread sklearn/metrics/_ranking.py Outdated
Comment on lines +1040 to +1042
# Micro-unit fixed point: enough resolution for typical weights while
# keeping scaled totals inside int64 for very large n. int64 is assumed
# available on float32-only devices of interest (e.g. torch MPS).

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks. This could be mentioned in the documentation so that users would know that granularity of the weights is limited.

# available on float32-only devices of interest (e.g. torch MPS).
scale = 1_000_000
y_true_i = xp.astype(y_true, xp.int64)
w_scaled = xp.astype(xp.round(weight * scale), xp.int64)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

weight = weight / weight.sum() gives mean of 1/n that get rounded to zero so
round(weight * 1e6) returns zero for any n>4e6.

I measured fo U(0,1) weights only 75% survive at n=1e6 and at n=1e7 tps[-1] is 0 and roc_curve returns a NaN where main returns 7.4e-4 relative error.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@Fazel94 I don't understand what you mean. This code does not divide by weight.sum(). Could you give a reproducer where this PR fails?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think I got what you meant. If the weights are normalized ahead of time (e.g. by the user, then using the 1e7 fixed scale can fail).

I iterated on this problem with an LLM in ogrisel#24. Let me merge this iteration into this PR as I actually think it's an improvement.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes that is exactly what i meant

Choose the micro-unit scale from the mean sample weight (capped by int64
headroom via n * max(weight)) so pre-normalized or uniformly tiny weights
do not round to zero under a fixed 1e6 scale, while O(1) weights keep the
previous micro-unit resolution. Extend the float32-only large-n regression
test to cover those failure modes.

Co-authored-by: Olivier Grisel <olivier.grisel@ensta.org>
…-4d84

FIX adaptive fixed-point scale for weighted float32-only cumsums
@github-actions github-actions Bot added the CI:Linter failure The linter CI is failing on this PR label Aug 28, 2026
@github-actions

Copy link
Copy Markdown

❌ Linting issues

This PR is introducing linting issues. Here's a summary of the issues. Note that you can avoid having linting issues by enabling pre-commit hooks. Instructions to enable them can be found here.

You can see the details of the linting issues under the lint job here


ruff check

ruff detected issues. Please run ruff check --fix --output-format=full locally, fix the remaining issues, and push the changes. Here you can see the detected issues. Note that the installed ruff version is ruff=0.12.2.

Details

sklearn/metrics/_ranking.py:1056:25: RUF046 Value being cast to `int` is already an integer
     |
1054 |             max_scaled_total = 0.9 * (2**63 - 1)
1055 |             scale_max = max_scaled_total / (n_samples * weight_max)
1056 |             scale = max(int(round(min(scale_res, scale_max))), 1)
     |                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ RUF046
1057 |         else:
1058 |             scale = target_units
     |
     = help: Remove unnecessary `int` call

Found 1 error.
No fixes available (1 hidden fix can be enabled with the `--unsafe-fixes` option).

Generated for commit: c6f51d5. Link to the linter CI: here

@Fazel94

Fazel94 commented Aug 28, 2026

Copy link
Copy Markdown
Contributor

In my tests your adaptive scale idea works great in all but one easily rectifiable case.

The case goes as w = 1e-8 on the top 10% by score, 1.0 on the rest, n = 1e7:

  scale=1.11e+06  round(1e-8*scale)=0
  tps[500000]  exact 0.0045005   #34827v2 0

tps would be wrong for the starting ones.
This can be fixed by consuming more of 64 bits you got by something in the line of

scale = max(int(0.9 * (2**63 - 1) / (n_samples * weight_max)), 1)

Generally summing in descending order is numerically problematic, the idea to use int64 is quiet good, the classic way of dealing with it is using compensated summing[1] that I tested a vectorized version[2] of it head to head to this version and the current pr is better.

def compensated_cumsum(x):
    c = np.cumsum(x, dtype=f32)
    a = np.concatenate([np.zeros(1, dtype=f32), c[:-1]])
    s = a + x
    bb = s - a
    e = (a - (s - bb)) + (x - bb)
    d = (s - c) + e
    return c + np.cumsum(d, dtype=f32)

[1]: Higham, Nicholas J. 4.6 In Accuracy and Stability of Numerical Algorithms, 2nd ed.
[2]: Knuth D. ch 4.2.2 The Art of Computer Programming, Vol. 2

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

Status: In Progress

Development

Successfully merging this pull request may close these issues.

4 participants