Make csr_row_norms support fused types

yenchenlin · yenchenlin · commit 65e91f715cc7 · 2016-05-19T08:49:17.000+08:00
diff --git a/sklearn/utils/sparsefuncs_fast.pyx b/sklearn/utils/sparsefuncs_fast.pyx
@@ -23,24 +23,29 @@ ctypedef np.float64_t DOUBLE
 
 def csr_row_norms(X):
     """L2 norm of each row in CSR matrix X."""
+    if X.dtype != np.float32:
+        X = X.astype(np.float64)
+    return _csr_row_norms(X.data, X.shape, X.indices, X.indptr)
+
+
+def _csr_row_norms(np.ndarray[floating, ndim=1, mode="c"] X_data,
+                   shape,
+                   np.ndarray[int, ndim=1, mode="c"] X_indices,
+                   np.ndarray[int, ndim=1, mode="c"] X_indptr):
     cdef:
-        unsigned int n_samples = X.shape[0]
-        unsigned int n_features = X.shape[1]
+        unsigned int n_samples = shape[0]
+        unsigned int n_features = shape[1]
         np.ndarray[DOUBLE, ndim=1, mode="c"] norms
-        np.ndarray[DOUBLE, ndim=1, mode="c"] data
-        np.ndarray[int, ndim=1, mode="c"] indices = X.indices
-        np.ndarray[int, ndim=1, mode="c"] indptr = X.indptr
 
         np.npy_intp i, j
         double sum_
 
     norms = np.zeros(n_samples, dtype=np.float64)
-    data = np.asarray(X.data, dtype=np.float64)     # might copy!
 
     for i in range(n_samples):
         sum_ = 0.0
-        for j in range(indptr[i], indptr[i + 1]):
-            sum_ += data[j] * data[j]
+        for j in range(X_indptr[i], X_indptr[i + 1]):
+            sum_ += X_data[j] * X_data[j]
         norms[i] = sum_
 
     return norms
