Forward-merge release/26.08 into main

python/cuml/cuml/linear_model/logistic_regression_mg.pyx

@@ -1,5 +1,5 @@
 #
-# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 import cupy as cp
@@ -211,6 +211,16 @@ class LogisticRegressionMG(LogisticRegression):
         classes.sort()
         cdef int n_classes = len(classes)
 
+        # QN's classification losses expect labels to be dense class indices
+        # in [0, n_classes). Keep the original values in ``classes`` for
+        # predictions, but encode the solver-facing labels using the global
+        # class ordering shared by all ranks.
+        y = cp.searchsorted(cp.asarray(classes, dtype=X.dtype), y).astype(
+            X.dtype, copy=False
+        )
+        opg.free_data_t(y_ptr, X.dtype)
+        y_ptr = opg.build_data_t([y])
+
         # Validate and initialize parameters
         l1_strength, l2_strength = self._get_l1_l2_strength()
         cdef qn_params params

python/cuml/tests/dask/test_dask_logistic_regression.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 
@@ -56,6 +56,9 @@ def _prep_training_data_sparse(c, X_train, y_train, partitions_per_worker):
 
     def cal_chunks(dataset, n_partitions):
         n_samples = dataset.shape[0]
+        # Avoid zero-row partitions when a small test dataset has fewer rows
+        # than the number of partitions requested by the available workers.
+        n_partitions = min(n_partitions, max(n_samples, 1))
         n_samples_per_part = int(n_samples / n_partitions)
         chunk_sizes = [n_samples_per_part] * n_partitions
         samples_last_row = n_samples - (
@@ -315,13 +318,16 @@ def test_noreg(fit_intercept, client):
 
 
 def test_n_classes_small(client):
-    def assert_small(X, y, n_classes):
+    def assert_small(X, y, n_classes, check_predictions=False):
         X_df, y_df = _prep_training_data(client, X, y, partitions_per_worker=1)
         from cuml.dask.linear_model import LogisticRegression as cumlLBFGS_dask
 
         lr = cumlLBFGS_dask()
         lr.fit(X_df, y_df)
         assert len(lr.classes_) == n_classes
+        if check_predictions:
+            predictions = lr.predict(X_df, delayed=True).compute().to_numpy()
+            assert np.array_equal(predictions, y)
         return lr
 
     X = np.array([(1, 2), (1, 3)], np.float32)
@@ -341,7 +347,7 @@ def assert_small(X, y, n_classes):
 
     X = np.array([(1, 2), (1, 3), (1, 2.5)], np.float32)
     y = np.array([10.0, 50.0, 20.0], np.float32)
-    lr = assert_small(X=X, y=y, n_classes=3)
+    lr = assert_small(X=X, y=y, n_classes=3, check_predictions=True)
     assert np.array_equal(
         lr.classes_, np.array([10.0, 20.0, 50.0], np.float32)
     )