diff --git a/python/cuml/cuml/_thirdparty/sklearn/preprocessing/_column_transformer.py b/python/cuml/cuml/_thirdparty/sklearn/preprocessing/_column_transformer.py index 466bd5da3d..8357d8ddad 100644 --- a/python/cuml/cuml/_thirdparty/sklearn/preprocessing/_column_transformer.py +++ b/python/cuml/cuml/_thirdparty/sklearn/preprocessing/_column_transformer.py @@ -738,6 +738,15 @@ def _validate_remainder(self, X): if hasattr(X, 'columns'): self._df_columns = X.columns + if hasattr(X, "shape"): + self._n_features = X.shape[1] + elif hasattr(X, "__len__") and len(X) > 0 and hasattr(X[0], "__len__"): + self._n_features = len(X[0]) + else: + raise TypeError( + "Input 'X' must be a 2D array, dataframe, or rectangular nested sequence." + ) + self._n_features = X.shape[1] cols = [] for columns in self._columns: diff --git a/python/cuml/tests/test_compose.py b/python/cuml/tests/test_compose.py index 665ad12b8a..d50f7e50f6 100644 --- a/python/cuml/tests/test_compose.py +++ b/python/cuml/tests/test_compose.py @@ -383,3 +383,15 @@ def test_column_transform_properly_handles_sub_output_type(): ] ).fit(df) transformer.transform(df) + +def test_make_column_transformer_list_input(): + from cuml.compose import make_column_transformer + from cuml.preprocessing import StandardScaler + import numpy as np + + a = [[1, 2], [3, 4], [5, 6]] + transformer = make_column_transformer((StandardScaler(), [0])) + res = transformer.fit_transform(a) + + expected = np.array([[-1.22474487], [0.0], [1.22474487]]) + np.testing.assert_allclose(res, expected, rtol=1e-5, atol=1e-5)