set dtype in dist_fix_point_cd

scikit-learn-contrib · Nov 8, 2023 · c336643 · c336643
1 parent 87fbfcd
commit c336643
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/skglm/solvers/common.py b/skglm/solvers/common.py
@@ -31,7 +31,7 @@ def dist_fix_point_cd(w, grad_ws, lipschitz, datafit, penalty, ws):
     dist : array, shape (n_features,)
         Violation score for every feature.
     """
-    dist = np.zeros(ws.shape[0])
+    dist = np.zeros(ws.shape[0], dtype=w.dtype)
 
     for idx, j in enumerate(ws):
         if lipschitz[j] == 0.:

diff --git a/skglm/solvers/prox_newton.py b/skglm/solvers/prox_newton.py
@@ -194,7 +194,7 @@ def solve(self, X, y, datafit, penalty, w_init=None, Xw_init=None):
         return w, np.asarray(p_objs_out), stop_crit
 
 
-# @njit
+@njit
 def _descent_direction(X, y, w_epoch, Xw_epoch, fit_intercept, grad_ws, datafit,
                        penalty, ws, tol, ws_strategy):
     # Given: