FEA Allow linear estimators to accept any dtype by converting to floa…

…t32 by default
rapidsai · May 14, 2024 · d8576f6 · d8576f6
1 parent 68d4336
commit d8576f6
Show file tree

Hide file tree

Showing 5 changed files with 17 additions and 3 deletions.
diff --git a/python/cuml/linear_model/linear_regression.pyx b/python/cuml/linear_model/linear_regression.pyx
@@ -331,6 +331,8 @@ class LinearRegression(LinearPredictMixin,
 
         X_m, n_rows, self.n_features_in_, self.dtype = \
             input_to_cuml_array(X,
+                                convert_to_dtype=(np.float32 if convert_dtype
+                                                  else None),
                                 check_dtype=[np.float32, np.float64],
                                 deepcopy=need_explicit_copy)
         _X_ptr = X_m.ptr

diff --git a/python/cuml/linear_model/ridge.pyx b/python/cuml/linear_model/ridge.pyx
@@ -252,6 +252,8 @@ class Ridge(UniversalBase,
         cdef uintptr_t _X_ptr, _y_ptr, _sample_weight_ptr
         X_m, n_rows, self.n_features_in_, self.dtype = \
             input_to_cuml_array(X, deepcopy=True,
+                                convert_to_dtype=(np.float32 if convert_dtype
+                                                  else None),
                                 check_dtype=[np.float32, np.float64])
         _X_ptr = X_m.ptr
         self.feature_names_in_ = X_m.index

diff --git a/python/cuml/solvers/cd.pyx b/python/cuml/solvers/cd.pyx
@@ -227,7 +227,10 @@ class CD(Base,
         """
         cdef uintptr_t sample_weight_ptr
         X_m, n_rows, self.n_cols, self.dtype = \
-            input_to_cuml_array(X, check_dtype=[np.float32, np.float64])
+            input_to_cuml_array(X,
+                                convert_to_dtype=(np.float32 if convert_dtype
+                                                  else None),
+                                check_dtype=[np.float32, np.float64])
 
         y_m, *_ = \
             input_to_cuml_array(y, check_dtype=self.dtype,

diff --git a/python/cuml/solvers/qn.pyx b/python/cuml/solvers/qn.pyx
@@ -470,7 +470,11 @@ class QN(Base,
         # Handle dense inputs
         else:
             X_m, n_rows, self.n_cols, self.dtype = input_to_cuml_array(
-                X, check_dtype=[np.float32, np.float64], order='K'
+                X,
+                convert_to_dtype=(np.float32 if convert_dtype
+                                  else None),
+                check_dtype=[np.float32, np.float64],
+                order='K'
             )
 
         y_m, _, _, _ = input_to_cuml_array(

diff --git a/python/cuml/solvers/sgd.pyx b/python/cuml/solvers/sgd.pyx
@@ -320,7 +320,10 @@ class SGD(Base,
 
         """
         X_m, n_rows, self.n_cols, self.dtype = \
-            input_to_cuml_array(X, check_dtype=[np.float32, np.float64])
+            input_to_cuml_array(X,
+                                convert_to_dtype=(np.float32 if convert_dtype
+                                                  else None),
+                                check_dtype=[np.float32, np.float64])
 
         y_m, _, _, _ = \
             input_to_cuml_array(y, check_dtype=self.dtype,