ENH csr_row_norms optimization (#24426)

ArturoAmorQ · jjerphan · glemaitre · web-flow · commit f965fcc0634e · 2023-01-27T15:17:02.000+01:00
Co-authored-by: Julien Jerphanion &lt;git@jjerphan.xyz&gt;
Co-authored-by: Guillaume Lemaitre &lt;g.lemaitre58@gmail.com&gt;
Co-authored-by: Jérémie du Boisberranger &lt;jeremiedbb@users.noreply.github.com&gt;
diff --git a/sklearn/utils/sparsefuncs_fast.pyx b/sklearn/utils/sparsefuncs_fast.pyx
@@ -24,32 +24,32 @@ ctypedef cnp.float64_t DOUBLE
 
 
 def csr_row_norms(X):
-    """L2 norm of each row in CSR matrix X."""
+    """Squared L2 norm of each row in CSR matrix X."""
     if X.dtype not in [np.float32, np.float64]:
         X = X.astype(np.float64)
-    return _csr_row_norms(X.data, X.shape, X.indices, X.indptr)
+    return _csr_row_norms(X.data, X.indices, X.indptr)
 
 
-def _csr_row_norms(cnp.ndarray[floating, ndim=1, mode="c"] X_data,
-                   shape,
-                   cnp.ndarray[integral, ndim=1, mode="c"] X_indices,
-                   cnp.ndarray[integral, ndim=1, mode="c"] X_indptr):
+def _csr_row_norms(
+    const floating[::1] X_data,
+    const integral[::1] X_indices,
+    const integral[::1] X_indptr,
+):
     cdef:
-        unsigned long long n_samples = shape[0]
-        unsigned long long i
-        integral j
+        integral n_samples = X_indptr.shape[0] - 1
+        integral i, j
         double sum_
 
-    norms = np.empty(n_samples, dtype=X_data.dtype)
-    cdef floating[::1] norms_view = norms
+    dtype = np.float32 if floating is float else np.float64
 
-    for i in range(n_samples):
-        sum_ = 0.0
-        for j in range(X_indptr[i], X_indptr[i + 1]):
-            sum_ += X_data[j] * X_data[j]
-        norms_view[i] = sum_
+    cdef floating[::1] norms = np.zeros(n_samples, dtype=dtype)
+
+    with nogil:
+        for i in range(n_samples):
+            for j in range(X_indptr[i], X_indptr[i + 1]):
+                norms[i] += X_data[j] * X_data[j]
 
-    return norms
+    return np.asarray(norms)
 
 
 def csr_mean_variance_axis0(X, weights=None, return_sum_weights=False):
