fix(ensemble, metrics): compute chunk sizes and refactor r2_score with nested da.where to prevent broadcasting errors

Author: wietzesuijker

Diff for: dask_ml/ensemble/_blockwise.py

@@ -54,6 +54,7 @@ def _predict(self, X):
             dtype = "float64"
 
         if isinstance(X, da.Array):
+            X = X.compute_chunk_sizes()  # Ensure integer chunks to avoid broadcasting errors.
             chunks = (X.chunks[0], len(self.estimators_))
             combined = X.map_blocks(
                 _predict_stack,
@@ -174,6 +175,7 @@ def _predict_proba(self, X):
 
     def _collect_probas(self, X):
         if isinstance(X, da.Array):
+            X = X.compute_chunk_sizes()  # Ensure integer chunks to avoid broadcasting errors.
             chunks = (len(self.estimators_), X.chunks[0], len(self.classes_))
             meta = np.array([], dtype="float64")
             # (n_estimators, len(X), n_classes)

Diff for: dask_ml/metrics/regression.py

@@ -155,25 +155,48 @@ def r2_score(
     multioutput: Optional[str] = "uniform_average",
     compute: bool = True,
 ) -> ArrayLike:
+    """
+    Compute the R² score for regression.
+
+    This function calculates the coefficient of determination using residual
+    and total sums of squares. It employs dask.array.where to gracefully handle
+    unknown dimensions without in-place assignment.
+
+    Parameters
+    ----------
+    y_true : ArrayLike
+        True target values.
+    y_pred : ArrayLike
+        Predicted target values.
+    sample_weight : Optional[ArrayLike], default=None
+        Weights for samples.
+    multioutput : Optional[str], default="uniform_average"
+        Method to aggregate multiple outputs.
+    compute : bool, default=True
+        If True, return the computed result; else, return a Dask array.
+
+    Returns
+    -------
+    result : ArrayLike
+        The R² score (scalar/NumPy array if computed, or a Dask array otherwise).
+    """
     _check_sample_weight(sample_weight)
     _, y_true, y_pred, _ = _check_reg_targets(y_true, y_pred, multioutput)
     weight = 1.0
 
+    # Compute residual and total sums of squares.
     numerator = (weight * (y_true - y_pred) ** 2).sum(axis=0, dtype="f8")
     denominator = (weight * (y_true - y_true.mean(axis=0)) ** 2).sum(axis=0, dtype="f8")
 
-    nonzero_denominator = denominator != 0
-    nonzero_numerator = numerator != 0
-    valid_score = nonzero_denominator & nonzero_numerator
-    output_chunks = getattr(y_true, "chunks", [None, None])[1]
-    output_scores = da.ones([y_true.shape[1]], chunks=output_chunks)
-    with np.errstate(all="ignore"):
-        output_scores[valid_score] = 1 - (
-            numerator[valid_score] / denominator[valid_score]
-        )
-        output_scores[nonzero_numerator & ~nonzero_denominator] = 0.0
-
-    result = output_scores.mean(axis=0)
+    # Determine R²: 1.0 for perfect predictions, 1 - numerator/denom when valid,
+    # and 0.0 if denominator is zero.
+    score = da.where(
+        numerator == 0,
+        1.0,
+        da.where(denominator != 0, 1 - numerator / denominator, 0.0)
+    )
+
+    result = score.mean(axis=0)
     if compute:
         result = result.compute()
     return result