fix: broadcast errors using lazy n_samples and da.where in r2_score

[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

• Track sample count in fit() using X.shape[0].
• Use n_samples to derive the rechunking block size so test data aligns one-to-one with training blocks, preventing broadcast mismatches.
• Refactor r2_score() to use da.where() for correct broadcasting.

Testing
Tests added to ensure:

• BlockwiseVotingRegressor supports various chunking patterns.
• 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.