test: r2_score and blockwise voting with different chunk patterns

Author: wietzesuijker

Diff for: tests/ensemble/test_blockwise.py

@@ -212,6 +212,25 @@ def test_fit_array(self):
         # score2 = est.score(X3, y3)
         # assert score == score2
 
+    def test_predict_with_different_chunks(self):
+        """Test prediction with different chunking patterns."""
+        # Train with one chunking pattern
+        X, y = dask_ml.datasets.make_regression(n_features=20, chunks=25)
+        est = dask_ml.ensemble.BlockwiseVotingRegressor(
+            sklearn.linear_model.LinearRegression(),
+        )
+        est.fit(X, y)
+        # Predict with a different chunking pattern
+        X_test, y_test = dask_ml.datasets.make_regression(n_features=20, chunks=20)
+
+        result = est.predict(X_test)
+        assert result.dtype == np.dtype("float64")
+        assert result.shape == y_test.shape
+
+        # Also test scoring
+        score = est.score(X_test, y_test)
+        assert isinstance(score, float)
+
     def test_fit_frame(self):
         X, y = dask_ml.datasets.make_regression(n_features=20, chunks=25)
         X = dd.from_dask_array(X)

Diff for: tests/metrics/test_regression.py

@@ -116,3 +116,19 @@ def test_regression_metrics_do_not_support_weighted_multioutput(metric_pairs):
 
     with pytest.raises((NotImplementedError, ValueError), match=error_msg):
         _ = m1(a, b, multioutput=weights)
+
+
+def test_r2_score_with_different_chunk_patterns():
+    """Test r2_score with different chunking configurations."""
+    # Create arrays with compatible but different chunk patterns
+    a = da.random.uniform(size=(100,), chunks=25)  # 4 chunks
+    b = da.random.uniform(size=(100,), chunks=20)  # 5 chunks
+    result = dask_ml.metrics.r2_score(a, b)
+    assert isinstance(result, float)
+    # Create arrays with different chunk patterns
+    a_multi = da.random.uniform(size=(100, 3), chunks=(25, 3))  # 4 chunks
+    b_multi = da.random.uniform(size=(100, 3), chunks=(20, 3))  # 5 chunks
+    result_multi = dask_ml.metrics.r2_score(
+        a_multi, b_multi, multioutput="uniform_average"
+    )
+    assert isinstance(result_multi, float)