diff --git a/sklearn/utils/sparsefuncs_fast.pyx b/sklearn/utils/sparsefuncs_fast.pyx index fdf59957f0afe..1d3eaf98916eb 100644 --- a/sklearn/utils/sparsefuncs_fast.pyx +++ b/sklearn/utils/sparsefuncs_fast.pyx @@ -11,6 +11,7 @@ cimport numpy as np import numpy as np import scipy.sparse as sp cimport cython +from cython cimport floating np.import_array() @@ -360,11 +361,11 @@ cdef void add_row_csr(np.ndarray[np.float64_t, ndim=1] data, def assign_rows_csr(X, np.ndarray[np.npy_intp, ndim=1] X_rows, np.ndarray[np.npy_intp, ndim=1] out_rows, - np.ndarray[np.float64_t, ndim=2, mode="c"] out): + np.ndarray[floating, ndim=2, mode="c"] out): """Densify selected rows of a CSR matrix into a preallocated array. Like out[out_rows] = X[X_rows].toarray() but without copying. - Only supported for dtype=np.float64. + No-copy supported for both dtype=np.float32 and dtype=np.float64. Parameters ---------- @@ -378,7 +379,7 @@ def assign_rows_csr(X, # but int is what scipy.sparse uses. int i, ind, j np.npy_intp rX - np.ndarray[DOUBLE, ndim=1] data = X.data + np.ndarray[floating, ndim=1] data = X.data np.ndarray[int, ndim=1] indices = X.indices, indptr = X.indptr if X_rows.shape[0] != out_rows.shape[0]: diff --git a/sklearn/utils/tests/test_sparsefuncs.py b/sklearn/utils/tests/test_sparsefuncs.py index 072144d7764e9..06c7ba4a342b8 100644 --- a/sklearn/utils/tests/test_sparsefuncs.py +++ b/sklearn/utils/tests/test_sparsefuncs.py @@ -179,20 +179,21 @@ def test_mean_variance_illegal_axis(): def test_densify_rows(): - X = sp.csr_matrix([[0, 3, 0], - [2, 4, 0], - [0, 0, 0], - [9, 8, 7], - [4, 0, 5]], dtype=np.float64) - X_rows = np.array([0, 2, 3], dtype=np.intp) - out = np.ones((6, X.shape[1]), dtype=np.float64) - out_rows = np.array([1, 3, 4], dtype=np.intp) - - expect = np.ones_like(out) - expect[out_rows] = X[X_rows, :].toarray() - - assign_rows_csr(X, X_rows, out_rows, out) - assert_array_equal(out, expect) + for dtype in (np.float32, np.float64): + X = sp.csr_matrix([[0, 3, 0], + [2, 4, 0], + [0, 0, 0], + [9, 8, 7], + [4, 0, 5]], dtype=dtype) + X_rows = np.array([0, 2, 3], dtype=np.intp) + out = np.ones((6, X.shape[1]), dtype=dtype) + out_rows = np.array([1, 3, 4], dtype=np.intp) + + expect = np.ones_like(out) + expect[out_rows] = X[X_rows, :].toarray() + + assign_rows_csr(X, X_rows, out_rows, out) + assert_array_equal(out, expect) def test_inplace_column_scale():