Refactor API context managers

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

@@ -296,7 +296,7 @@ def _list_indexing(X, key, key_dtype):
 
 
 def _transform_one(transformer, X, y, weight, **fit_params):
-    res = transformer.transform(X).to_output('cupy')
+    res = transformer.transform(X).to_output("array")
     # if we have a weight for this transformer, multiply output
     if weight is None:
         return res
@@ -316,16 +316,15 @@ def _fit_transform_one(transformer,
     be multiplied by ``weight``.
     """
     with _print_elapsed_time(message_clsname, message):
-        with cuml.using_output_type("cupy"):
-            transformer.accept_sparse = True
-            if hasattr(transformer, 'fit_transform'):
-                res = transformer.fit_transform(X, y, **fit_params)
-            else:
-                res = transformer.fit(X, y, **fit_params).transform(X)
+        transformer.accept_sparse = True
+        if hasattr(transformer, 'fit_transform'):
+            res = transformer.fit_transform(X, y, **fit_params)
+        else:
+            res = transformer.fit(X, y, **fit_params).transform(X)
 
     if weight is None:
         return res, transformer
-    return res * weight, transformer
+    return res.to_output("array") * weight, transformer
 
 
 def _name_estimators(estimators):
@@ -899,6 +898,7 @@ def fit_transform(self, X, y=None) -> SparseCumlArray:
             return np.zeros((X.shape[0], 0))
 
         Xs, transformers = zip(*result)
+        Xs = [x.to_output("array") if hasattr(x, "to_output") else x for x in Xs]
 
         # determine if concatenated output will be sparse or not
         if any(issparse(X) for X in Xs):

python/cuml/common/array_descriptor.py

@@ -16,7 +16,7 @@
 
 from dataclasses import dataclass, field
 from cuml.internals.array import CumlArray
-import cuml
+from cuml.internals.global_settings import GlobalSettings
 from cuml.internals.input_utils import (
     input_to_cuml_array, determine_array_type
 )
@@ -114,7 +114,7 @@ def __get__(self, instance, owner):
         assert len(existing.values) > 0
 
         # Get the global output type
-        output_type = cuml.global_settings.output_type
+        output_type = GlobalSettings().output_type
 
         # First, determine if we need to call to_output at all
         if (output_type == "mirror"):

python/cuml/decomposition/pca.pyx

@@ -497,7 +497,6 @@ class PCA(UniversalBase,
                                        'type': 'dense_sparse',
                                        'description': 'Transformed values',
                                        'shape': '(n_samples, n_components)'})
-    @cuml.internals.api_base_return_array_skipall
     @enable_device_interop
     def fit_transform(self, X, y=None) -> CumlArray:
         """

python/cuml/feature_extraction/_tfidf.py

@@ -188,7 +188,7 @@ def fit(self, X) -> "TfidfTransformer":
 
         return self
 
-    @cuml.internals.api_base_return_any_skipall
+    @cuml.internals.api_base_return_array()
     def transform(self, X, copy=True):
         """Transform a count matrix to a tf or tf-idf representation
 
@@ -238,7 +238,7 @@ def transform(self, X, copy=True):
 
         return X
 
-    @cuml.internals.api_base_return_any_skipall
+    @cuml.internals.api_base_return_array()
     def fit_transform(self, X, copy=True):
         """
         Fit TfidfTransformer to X, then transform X.

python/cuml/internals/__init__.py

@@ -21,20 +21,18 @@
 from cuml.internals.api_decorators import (
     _deprecate_pos_args,
     api_base_fit_transform,
-    api_base_return_any_skipall,
     api_base_return_any,
-    api_base_return_array_skipall,
+    api_base_return_any_skipall,
     api_base_return_array,
-    api_base_return_generic_skipall,
+    api_base_return_array_skipall,
     api_base_return_generic,
+    api_base_return_generic_skipall,
     api_base_return_sparse_array,
     api_return_any,
     api_return_array,
     api_return_generic,
     api_return_sparse_array,
     exit_internal_api,
-)
-from cuml.internals.api_context_managers import (
     in_internal_api,
     set_api_output_dtype,
     set_api_output_type,

python/cuml/internals/api_context.py

@@ -0,0 +1,25 @@
+#
+# Copyright (c) 2023, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+from dataclasses import dataclass
+
+
+@dataclass
+class ApiContext:
+
+    stack_level: int = 0
+    previous_output_type = None
+    output_type = None
+    output_dtype = None