Simplify type reflection implementation

docs/source/api.rst

@@ -15,8 +15,9 @@ Module Configuration
 Output Data Type Configuration
 ------------------------------
 
- .. autofunction:: cuml.internals.memory_utils.set_global_output_type
- .. autofunction:: cuml.internals.memory_utils.using_output_type
+.. autofunction:: cuml.set_global_output_type
+
+.. autofunction:: cuml.using_output_type
 
 .. _verbosity-levels:
 

docs/source/cuml_intro.rst

@@ -77,10 +77,8 @@ fit a model with a NumPy array, the ``model.coef_`` property
 containing fitted coefficients will also be a NumPy array. If you fit
 a model using cuDF's GPU-based DataFrame and Series objects, the
 model's output properties will be cuDF objects. You can always
-override this behavior and select a default datatype with the
-`memory_utils.set_global_output_type
-<https://docs.rapids.ai/api/cuml/nightly/api.html#datatype-configuration>`_
-function.
+override this behavior and select a default datatype with
+:func:`cuml.set_global_output_type`.
 
 The `RAPIDS Configurable Input and Output Types
 <https://medium.com/@dantegd/e719d72c135b>`_ blog post goes into much

python/cuml/cuml/__init__.py

@@ -42,10 +42,7 @@
     GlobalSettings,
     _global_settings_data,
 )
-from cuml.internals.memory_utils import (
-    set_global_output_type,
-    using_output_type,
-)
+from cuml.internals.outputs import set_global_output_type, using_output_type
 from cuml.kernel_ridge.kernel_ridge import KernelRidge
 from cuml.linear_model.elastic_net import ElasticNet
 from cuml.linear_model.lasso import Lasso

python/cuml/cuml/_thirdparty/sklearn/preprocessing/_column_transformer.py

@@ -833,6 +833,7 @@ def _fit_transform(self, X, y, func, fitted=False):
             else:
                 raise
 
+    @cuml.internals.reflect
     def fit(self, X, y=None) -> "ColumnTransformer":
         """Fit all transformers using X.
 
@@ -856,6 +857,7 @@ def fit(self, X, y=None) -> "ColumnTransformer":
         self.fit_transform(X, y=y)
         return self
 
+    @cuml.internals.reflect(reset=True)
     def fit_transform(self, X, y=None) -> SparseCumlArray:
         """Fit all transformers, transform the data and concatenate results.
 
@@ -913,6 +915,7 @@ def fit_transform(self, X, y=None) -> SparseCumlArray:
 
         return self._hstack(list(Xs))
 
+    @cuml.internals.reflect
     def transform(self, X) -> SparseCumlArray:
         """Transform X separately by each transformer, concatenate results.
 

python/cuml/cuml/_thirdparty/sklearn/preprocessing/_data.py

@@ -42,10 +42,9 @@
 )
 
 from ....common.array_descriptor import CumlArrayDescriptor
-from ....internals import api_return_generic
 from ....internals.array import CumlArray
 from ....internals.array_sparse import SparseCumlArray
-from ....internals.memory_utils import using_output_type
+from ....internals.outputs import using_output_type, reflect
 from ....thirdparty_adapters import check_array
 from ....thirdparty_adapters.sparsefuncs_fast import (
     csr_polynomial_expansion,
@@ -102,7 +101,7 @@ def _handle_zeros_in_scale(scale, copy=True):
         return scale
 
 
-@api_return_generic(get_output_type=True)
+@reflect
 def scale(X, *, axis=0, with_mean=True, with_std=True, copy=True):
     """Standardize a dataset along any axis
 
@@ -330,6 +329,7 @@ def _get_param_names(cls):
             "copy"
         ]
 
+    @reflect(reset=True)
     def fit(self, X, y=None) -> "MinMaxScaler":
         """Compute the minimum and maximum to be used for later scaling.
 
@@ -352,6 +352,7 @@ def fit(self, X, y=None) -> "MinMaxScaler":
         self._reset()
         return self.partial_fit(X, y)
 
+    @reflect(reset=True)
     def partial_fit(self, X, y=None) -> "MinMaxScaler":
         """Online computation of min and max on X for later scaling.
 
@@ -402,6 +403,7 @@ def partial_fit(self, X, y=None) -> "MinMaxScaler":
         self.data_range_ = data_range
         return self
 
+    @reflect
     def transform(self, X) -> CumlArray:
         """Scale features of X according to feature_range.
 
@@ -425,6 +427,7 @@ def transform(self, X) -> CumlArray:
 
         return X
 
+    @reflect
     def inverse_transform(self, X) -> CumlArray:
         """Undo the scaling of X according to feature_range.
 
@@ -448,7 +451,7 @@ def inverse_transform(self, X) -> CumlArray:
         return X
 
 
-@api_return_generic(get_output_type=True)
+@reflect
 def minmax_scale(X, feature_range=(0, 1), *, axis=0, copy=True):
     """Transform features by scaling each feature to a given range.
 
@@ -655,6 +658,7 @@ def _get_param_names(cls):
             "copy"
         ]
 
+    @reflect(reset=True)
     def fit(self, X, y=None) -> "StandardScaler":
         """Compute the mean and std to be used for later scaling.
 
@@ -672,6 +676,7 @@ def fit(self, X, y=None) -> "StandardScaler":
         self._reset()
         return self.partial_fit(X, y)
 
+    @reflect(reset=True)
     def partial_fit(self, X, y=None) -> "StandardScaler":
         """
         Online computation of mean and std on X for later scaling.
@@ -792,6 +797,7 @@ def partial_fit(self, X, y=None) -> "StandardScaler":
 
         return self
 
+    @reflect
     def transform(self, X, copy=None) -> SparseCumlArray:
         """Perform standardization by centering and scaling
 
@@ -827,6 +833,7 @@ def transform(self, X, copy=None) -> SparseCumlArray:
 
         return X
 
+    @reflect
     def inverse_transform(self, X, copy=None) -> SparseCumlArray:
         """Scale back the data to the original representation
 
@@ -957,6 +964,7 @@ def _get_param_names(cls):
             "copy"
         ]
 
+    @reflect(reset=True)
     def fit(self, X, y=None) -> "MaxAbsScaler":
         """Compute the maximum absolute value to be used for later scaling.
 
@@ -971,6 +979,7 @@ def fit(self, X, y=None) -> "MaxAbsScaler":
         self._reset()
         return self.partial_fit(X, y)
 
+    @reflect(reset=True)
     def partial_fit(self, X, y=None) -> "MaxAbsScaler":
         """
         Online computation of max absolute value of X for later scaling.
@@ -1015,6 +1024,7 @@ def partial_fit(self, X, y=None) -> "MaxAbsScaler":
         self.scale_ = _handle_zeros_in_scale(max_abs)
         return self
 
+    @reflect
     def transform(self, X) -> SparseCumlArray:
         """Scale the data
 
@@ -1036,6 +1046,7 @@ def transform(self, X) -> SparseCumlArray:
 
         return X
 
+    @reflect
     def inverse_transform(self, X) -> SparseCumlArray:
         """Scale back the data to the original representation
 
@@ -1057,7 +1068,7 @@ def inverse_transform(self, X) -> SparseCumlArray:
         return X
 
 
-@api_return_generic(get_output_type=True)
+@reflect
 def maxabs_scale(X, *, axis=0, copy=True):
     """Scale each feature to the [-1, 1] range without breaking the sparsity.
 
@@ -1209,6 +1220,7 @@ def _get_param_names(cls):
             "copy"
         ]
 
+    @reflect(reset=True)
     def fit(self, X, y=None) -> "RobustScaler":
         """Compute the median and quantiles to be used for scaling.
 
@@ -1270,6 +1282,7 @@ def fit(self, X, y=None) -> "RobustScaler":
 
         return self
 
+    @reflect
     def transform(self, X) -> SparseCumlArray:
         """Center and scale the data.
 
@@ -1294,6 +1307,7 @@ def transform(self, X) -> SparseCumlArray:
                 X /= self.scale_
         return X
 
+    @reflect
     def inverse_transform(self, X) -> SparseCumlArray:
         """Scale back the data to the original representation
 
@@ -1319,7 +1333,7 @@ def inverse_transform(self, X) -> SparseCumlArray:
         return X
 
 
-@api_return_generic(get_output_type=True)
+@reflect
 def robust_scale(X, *, axis=0, with_centering=True, with_scaling=True,
                  quantile_range=(25.0, 75.0), copy=True):
     """
@@ -1529,6 +1543,7 @@ def get_feature_names(self, input_features=None):
             feature_names.append(name)
         return feature_names
 
+    @reflect(reset=True)
     def fit(self, X, y=None) -> "PolynomialFeatures":
         """
         Compute number of output features.
@@ -1552,6 +1567,7 @@ def fit(self, X, y=None) -> "PolynomialFeatures":
         self.n_output_features_ = sum(1 for _ in combinations)
         return self
 
+    @reflect
     def transform(self, X) -> SparseCumlArray:
         """Transform data to polynomial features
 
@@ -1679,7 +1695,7 @@ def transform(self, X) -> SparseCumlArray:
         return XP  # TODO keep order
 
 
-@api_return_generic(get_output_type=True)
+@reflect
 def normalize(X, norm='l2', *, axis=1, copy=True, return_norm=False):
     """Scale input vectors individually to unit norm (vector length).
 
@@ -1830,6 +1846,7 @@ def __init__(self, norm='l2', *, copy=True):
         self.norm = norm
         self.copy = copy
 
+    @reflect(reset=True)
     def fit(self, X, y=None) -> "Normalizer":
         """Do nothing and return the estimator unchanged
 
@@ -1843,6 +1860,7 @@ def fit(self, X, y=None) -> "Normalizer":
         self._validate_data(X, accept_sparse='csr')
         return self
 
+    @reflect
     def transform(self, X, copy=None) -> SparseCumlArray:
         """Scale each non zero row of X to unit norm
 
@@ -1859,7 +1877,7 @@ def transform(self, X, copy=None) -> SparseCumlArray:
         return normalize(X, norm=self.norm, axis=1, copy=copy)
 
 
-@api_return_generic(get_output_type=True)
+@reflect
 def binarize(X, *, threshold=0.0, copy=True):
     """Boolean thresholding of array-like or sparse matrix
 
@@ -1959,6 +1977,7 @@ def __init__(self, *, threshold=0.0, copy=True):
         self.threshold = threshold
         self.copy = copy
 
+    @reflect(reset=True)
     def fit(self, X, y=None) -> "Binarizer":
         """Do nothing and return the estimator unchanged
 
@@ -1972,6 +1991,7 @@ def fit(self, X, y=None) -> "Binarizer":
         self._validate_data(X, accept_sparse=['csr', 'csc'])
         return self
 
+    @reflect
     def transform(self, X, copy=None) -> SparseCumlArray:
         """Binarize each element of X
 
@@ -1988,7 +2008,7 @@ def transform(self, X, copy=None) -> SparseCumlArray:
         return binarize(X, threshold=self.threshold, copy=copy)
 
 
-@api_return_generic(get_output_type=True)
+@reflect
 def add_dummy_feature(X, value=1.0):
     """Augment dataset with an additional dummy feature.
 
@@ -2098,6 +2118,7 @@ def __init__(self):
         # Needed for backported inspect.signature compatibility with PyPy
         pass
 
+    @reflect(reset=True)
     def fit(self, K, y=None) -> 'KernelCenterer':
         """Fit KernelCenterer
 
@@ -2123,6 +2144,7 @@ def fit(self, K, y=None) -> 'KernelCenterer':
         self.K_fit_all_ = self.K_fit_rows_.sum() / n_samples
         return self
 
+    @reflect
     def transform(self, K, copy=True) -> CumlArray:
         """Center kernel matrix.
 
@@ -2355,6 +2377,7 @@ def _sparse_fit(self, X, random_state):
         # https://github.com/numpy/numpy/issues/14685
         self.quantiles_ = np.array(cpu_np.maximum.accumulate(self.quantiles_))
 
+    @reflect(reset=True)
     def fit(self, X, y=None) -> 'QuantileTransformer':
         """Compute the quantiles used for transforming.
 
@@ -2548,6 +2571,7 @@ def _transform(self, X, inverse=False):
 
         return X
 
+    @reflect
     def transform(self, X) -> SparseCumlArray:
         """Feature-wise transformation of the data.
 
@@ -2569,6 +2593,7 @@ def transform(self, X) -> SparseCumlArray:
 
         return self._transform(X, inverse=False)
 
+    @reflect
     def inverse_transform(self, X) -> SparseCumlArray:
         """Back-projection to the original space.
 
@@ -2795,6 +2820,7 @@ def _get_param_names(cls):
             "copy"
         ]
 
+    @reflect(reset=True)
     def fit(self, X, y=None) -> 'PowerTransformer':
         """Estimate the optimal parameter lambda for each feature.
 
@@ -2815,6 +2841,7 @@ def fit(self, X, y=None) -> 'PowerTransformer':
         self._fit(X, y=y, force_transform=False)
         return self
 
+    @reflect(reset=True)
     def fit_transform(self, X, y=None) -> CumlArray:
         return self._fit(X, y, force_transform=True)
 
@@ -2853,6 +2880,7 @@ def _fit(self, X, y=None, force_transform=False):
 
         return X
 
+    @reflect
     def transform(self, X) -> CumlArray:
         """Apply the power transform to each feature using the fitted lambdas.
 
@@ -2887,6 +2915,7 @@ def transform(self, X) -> CumlArray:
 
         return X
 
+    @reflect
     def inverse_transform(self, X) -> CumlArray:
         """Apply the inverse power transformation using the fitted lambdas.