diff --git a/skl2onnx/shape_calculators/function_transformer.py b/skl2onnx/shape_calculators/function_transformer.py index 4b2a0ade7..a475f9c2a 100644 --- a/skl2onnx/shape_calculators/function_transformer.py +++ b/skl2onnx/shape_calculators/function_transformer.py @@ -4,15 +4,15 @@ # license information. # -------------------------------------------------------------------------- +import copy from ..common._registration import register_shape_calculator -from ..common.data_types import FloatTensorType def calculate_sklearn_function_transformer_output_shapes(operator): - ''' + """ This operator is used only to merge columns in a pipeline. - Only id function is supported. - ''' + Only identity function is supported. + """ if operator.raw_operator.func is not None: raise RuntimeError("FunctionTransformer is not supported unless the " "transform function is None (= identity). " @@ -27,7 +27,8 @@ def calculate_sklearn_function_transformer_output_shapes(operator): C = 'None' break - operator.outputs[0].type = FloatTensorType([N, C]) + operator.outputs[0].type = copy.deepcopy(operator.inputs[0].type) + operator.outputs[0].type.shape = [N, C] register_shape_calculator('SklearnFunctionTransformer', diff --git a/tests/test_sklearn_function_transformer_converter.py b/tests/test_sklearn_function_transformer_converter.py index 59651c5bf..ee03169cf 100644 --- a/tests/test_sklearn_function_transformer_converter.py +++ b/tests/test_sklearn_function_transformer_converter.py @@ -2,8 +2,9 @@ Tests scikit-imputer converter. """ import unittest +import numpy as np import pandas -from sklearn.datasets import load_iris +from sklearn.datasets import load_digits, load_iris from sklearn.pipeline import Pipeline try: @@ -44,10 +45,10 @@ def convert_dataframe_schema(df, drop=None): inputs.append((k, t)) return inputs - data = load_iris() + data = load_digits() X = data.data[:, :2] y = data.target - data = pandas.DataFrame(X, columns=["X1", "X2"]) + data = pandas.DataFrame(X, columns=["X1", "X2"], dtype=np.int64) pipe = Pipeline(steps=[ (