FIX Remove Numba warning when fitting multi-task estimators (#93)

Co-authored-by: mathurinm <[email protected]>
Co-authored-by: Badr MOUFAD <[email protected]>

Author: 3 people

skglm/datafits/multi_task.py

@@ -68,7 +68,7 @@ def value(self, Y, W, XW):
     def gradient_j(self, X, Y, W, XW, j):
         """Gradient with respect to j-th coordinate of W."""
         n_samples = X.shape[0]
-        return (X[:, j:j+1].T @ XW - self.XtY[j, :]) / n_samples
+        return (X[:, j] @ XW - self.XtY[j, :]) / n_samples
 
     def gradient_j_sparse(self, X_data, X_indptr, X_indices, Y, XW, j):
         """Gradient with respect to j-th coordinate of W when X is sparse."""

skglm/solvers/multitask_bcd.py

@@ -369,8 +369,8 @@ def _bcd_epoch(X, Y, W, XW, datafit, penalty, ws):
             continue
         Xj = X[:, j]
         old_W_j = W[j, :].copy()  # copy is very important here
-        W[j:j+1, :] = penalty.prox_1feat(
-            W[j:j+1, :] - datafit.gradient_j(X, Y, W, XW, j) / lc[j],
+        W[j, :] = penalty.prox_1feat(
+            W[j, :] - datafit.gradient_j(X, Y, W, XW, j) / lc[j],
             1 / lc[j], j)
         if not np.all(W[j, :] == old_W_j):
             for k in range(n_tasks):