fix: prevent broadcasting errors in r2_score using da.where()

[wietzesuijker]

Closes #1012

First PR here. Curious to hear your feedback.

Problem
After updating to Dask 2025.2.0, tests fail with a ValueError due to changes in chunk size handling.

Solution
Refactor r2_score() to use da.where() for correct broadcasting.

Testing
Test added to ensure r2_score() works correctly with arrays that have different chunk configurations.

[TomAugspurger]

Thanks.

I'm not entirely sure what the best action is, but I think we ought to avoid anything that triggers computation unnecessarily, including len.

Can you say a bit more about getting n_samples is needed in blockwise?

[wietzesuijker]

Thanks @TomAugspurger. n_samples (now obtained via X.shape[0]) lets us determine the rechunking size without forcing computation. It splits the test data into one block per trained estimator, ensuring alignment and preventing broadcast errors. Combined with the da.where() update in r2_score, these changes maintain laziness and correct behavior with mismatched chunks.

[TomAugspurger]

I'm probably missing something, but why do we care that the size of the test dataset matches the size of the training dataset (_n_samples)? I'd expect us to just care that the number of samples in X_train and y_train to match, and separately that the number of samples in X_test, y_test match.

[wietzesuijker]

why do we care that the size of the test dataset matches the size of the training dataset (_n_samples)?

The goal is not for the test dataset to match the training dataset's overall size. The focus is ensuring each estimator, trained on a specific data block, receives a matching block from the test set. X.shape[0] is used (as n_samples) to compute the optimal test data block size, dividing the test set into blocks equal to the number of estimators. This aligns predictions and prevents broadcast errors, regardless of training and test dataset sizes.

[wietzesuijker]

@TomAugspurger I limited the PR to changes in dask_ml/metrics/regression.py. Merging this would unblock my use case and allow me to update dask. Thanks! :). (previous state)

[wietzesuijker]

Thanks for triggering the tests, Tom. Is there anything I can/should do to fix the failing runs? The errors seem unrelated and similar to other recent runs e.g. https://github.com/dask/dask-ml/actions/runs/14849752234/job/41691028477.

[TomAugspurger]

I'll take a look in the next few days.

[TomAugspurger]

Thanks @wietzesuijker. There's still one error in tests/preprocessing/test_data.py::TestQuantileTransformer::test_fit_transform_frame but the rest from #1012 are fixed. I'll fix that in another PR.