assert GPU CPU intercept_ equal when fit_intercept is false

rapidsai · Sep 18, 2023 · 18acacf · 18acacf
1 parent 9b1ee68
commit 18acacf
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/python/cuml/solvers/qn.pyx b/python/cuml/solvers/qn.pyx
@@ -930,8 +930,15 @@ class QN(Base,
 
         if self.fit_intercept:
             self.intercept_ = self._coef_[-1]
-        else:
+            return
+
+        _num_classes_dim, _ = self.coef_.shape
+        _num_classes = self.get_num_classes(_num_classes_dim)
+
+        if _num_classes == 2:
             self.intercept_ = CumlArray.zeros(shape=1)
+        else:
+            self.intercept_ = CumlArray.zeros(shape=_num_classes)
 
     def get_param_names(self):
         return super().get_param_names() + \

diff --git a/python/cuml/tests/test_linear_model.py b/python/cuml/tests/test_linear_model.py
@@ -561,6 +561,9 @@ def test_logistic_regression(
     )
     assert len(np.unique(cu_preds)) == len(np.unique(y_test))
 
+    if fit_intercept is False:
+        assert np.array_equal(culog.intercept_, sklog.intercept_)
+
 
 @given(
     dtype=floating_dtypes(sizes=(32, 64)),