Repository: onnx/sklearn-onnx
Pull request #1164 (2025-02-03 18:43:18)

QuantileTransformer converter
Fix issue https://github.com/onnx/sklearn-onnx/issues/608

xadupre: Thanks for doing this.
vaaven: Did it.

Patch:
diff --git a/skl2onnx/_supported_operators.py b/skl2onnx/_supported_operators.py
index 45ff50f..094f2cc 100644
--- a/skl2onnx/_supported_operators.py
+++ b/skl2onnx/_supported_operators.py
@@ -273,6 +273,12 @@ except ImportError:
     # Not available in scikit-learn < 0.20.0
     PowerTransformer = None
 
+try:
+    from sklearn.preprocessing import QuantileTransformer
+except ImportError:
+    # Not available in scikit-learn < 0.19.0
+    QuantileTransformer = None
+
 try:
     from sklearn.ensemble import (
         HistGradientBoostingClassifier,
@@ -447,6 +453,7 @@ def build_sklearn_operator_name_map():
             PolynomialFeatures,
             PowerTransformer,
             QuadraticDiscriminantAnalysis,
+            QuantileTransformer,
             RadiusNeighborsClassifier,
             RadiusNeighborsRegressor,
             RandomForestClassifier,
diff --git a/skl2onnx/operator_converters/__init__.py b/skl2onnx/operator_converters/__init__.py
index f5a7885..48672c5 100644
--- a/skl2onnx/operator_converters/__init__.py
+++ b/skl2onnx/operator_converters/__init__.py
@@ -50,6 +50,7 @@ from . import pipelines
 from . import polynomial_features
 from . import power_transformer
 from . import quadratic_discriminant_analysis
+from . import quantile_transformer
 from . import random_forest
 from . import random_projection
 from . import random_trees_embedding
@@ -118,6 +119,7 @@ __all__ = [
     polynomial_features,
     power_transformer,
     quadratic_discriminant_analysis,
+    quantile_transformer,
     random_forest,
     random_projection,
     random_trees_embedding,
diff --git a/skl2onnx/operator_converters/quantile_transformer.py b/skl2onnx/operator_converters/quantile_transformer.py
new file mode 100644
index 0000000..90c93ee
--- /dev/null
+++ b/skl2onnx/operator_converters/quantile_transformer.py
@@ -0,0 +1,263 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import numpy as np
+import math
+from ..common._registration import register_converter
+from ..common._topology import Scope, Operator
+from ..common._container import ModelComponentContainer
+from ..common.data_types import guess_numpy_type, guess_proto_type, TensorProto
+from ..algebra.onnx_ops import (
+    OnnxWhere,
+    OnnxIsNaN,
+    OnnxNot,
+    OnnxCast,
+    OnnxEqual,
+    OnnxNeg,
+    OnnxSub,
+    OnnxMul,
+    OnnxAdd,
+    OnnxDiv,
+    OnnxGather,
+    OnnxLessOrEqual,
+    OnnxReshape,
+    OnnxShape,
+    OnnxFlatten,
+    OnnxGreaterOrEqual,
+    OnnxOr,
+)
+
+
+def _next_power_of_two(n):
+    return 1 << (n).bit_length()
+
+
+def _pad_quantiles(quantiles, target_size, n_quantiles, mode="min"):
+    pad_before = target_size - n_quantiles + 1
+    if mode == "min":
+        pad_values = np.repeat(quantiles[0:1, :], pad_before, axis=0)
+    else:
+        pad_values = np.repeat(quantiles[-1:, :], pad_before, axis=0)
+    return np.concatenate([pad_values, quantiles], axis=0)
+
+
+def _pad_references(references, target_size, n_quantiles, mode="min"):
+    pad_before = target_size - n_quantiles + 1
+    if mode == "min":
+        pad_values = np.repeat([references[0]], pad_before)
+    else:
+        pad_values = np.repeat([references[-1]], pad_before)
+    return np.concatenate([pad_values, references])
+
+
+def convert_quantile_transformer(
+    scope: Scope, operator: Operator, container: ModelComponentContainer
+):
+    op = operator.raw_operator
+    op_in = operator.inputs[0]
+    op_out = operator.outputs[0].full_name
+    opv = container.target_opset
+
+    if op.output_distribution != "uniform":
+        raise RuntimeError(
+            "Conversion of QuantileTransformer with output_distribution=%r "
+            "is not supported." % op.output_distribution
+        )
+
+    dtype = guess_numpy_type(op_in.type)
+    if dtype != np.float64:
+        dtype = np.float32
+    proto_dtype = guess_proto_type(op_in.type)
+
+    quantiles = op.quantiles_.astype(dtype)
+    references = op.references_.astype(dtype)
+    n_quantiles, n_features = quantiles.shape
+
+    target_size = _next_power_of_two(n_quantiles)
+
+    # Correct padding with min/max values
+    quantiles_pad = _pad_quantiles(quantiles, target_size, n_quantiles, "min")
+    references_pad = _pad_references(references, target_size, n_quantiles, "min")
+
+    # Prepare reversed and negated quantiles and references
+    quantiles_neg_rev = -quantiles[::-1, :]
+    references_neg_rev = -references[::-1]
+
+    quantiles_neg_rev_pad = _pad_quantiles(
+        quantiles_neg_rev, target_size, n_quantiles, "min"
+    )
+    references_neg_rev_pad = _pad_references(
+        references_neg_rev, target_size, n_quantiles, "min"
+    )
+
+    # Prepare initializers
+    def prepare_initializer(arr, name, transpose=False):
+        if transpose:
+            arr = arr.T.flatten()
+        else:
+            arr = arr.flatten()
+        arr_name = scope.get_unique_variable_name(name)
+        container.add_initializer(arr_name, proto_dtype, arr.shape, arr)
+        return arr_name
+
+    quantiles_name = prepare_initializer(quantiles_pad, "quantiles", transpose=True)
+    references_name = prepare_initializer(
+        np.tile(references_pad, n_features), "references"
+    )
+    quantiles_neg_rev_name = prepare_initializer(
+        quantiles_neg_rev_pad, "quantiles_neg_rev", transpose=True
+    )
+    references_neg_rev_name = prepare_initializer(
+        np.tile(references_neg_rev_pad, n_features), "references_neg_rev"
+    )
+
+    is_nan = OnnxIsNaN(op_in, op_version=opv)
+    is_finite = OnnxNot(is_nan, op_version=opv)
+    X_clean = OnnxWhere(is_finite, op_in, op_in, op_version=opv)
+    X_clean.add_to(scope, container)
+
+    def create_interpolation(X_input, quantiles_init, references_init):
+        flat_input = OnnxFlatten(X_input, axis=0, op_version=opv)
+        flat_input.add_to(scope, container)
+
+        current_index = OnnxMul(
+            OnnxCast(X_input, to=np.int64, op_version=opv),
+            np.array(0, dtype=np.int64),
+            op_version=opv,
+        )
+        current_index = OnnxFlatten(
+            OnnxAdd(
+                current_index,
+                np.array(
+                    [i * (target_size + 1) for i in range(n_features)], dtype=np.int64
+                ),
+                op_version=opv,
+            ),
+            axis=0,
+            op_version=opv,
+        )
+        current_index.add_to(scope, container)
+
+        step = OnnxFlatten(
+            OnnxAdd(
+                OnnxMul(current_index, np.array(0, dtype=np.int64), op_version=opv),
+                np.array([target_size // 2], dtype=np.int64),
+                op_version=opv,
+            )
+        )
+        step.add_to(scope, container)
+
+        for _i in range(int(math.log2(target_size))):
+            current_index_temp = OnnxAdd(current_index, step, op_version=opv)
+            current_index_temp.add_to(scope, container)
+
+            q_current = OnnxGather(
+                quantiles_init, current_index_temp, axis=0, op_version=opv
+            )
+            q_current.add_to(scope, container)
+
+            mask = OnnxGreaterOrEqual(flat_input, q_current, op_version=opv)
+            mask.add_to(scope, container)
+
+            current_index = OnnxWhere(
+                mask, current_index_temp, current_index, op_version=opv
+            )
+            current_index.add_to(scope, container)
+
+            step = OnnxCast(
+                OnnxDiv(step, np.array(2, dtype=np.int64), op_version=opv),
+                to=TensorProto.INT64,
+                op_version=opv,
+            )
+            step.add_to(scope, container)
+
+        index_plus_one = OnnxAdd(
+            current_index, np.array(1, dtype=np.int64), op_version=opv
+        )
+        index_plus_one.add_to(scope, container)
+
+        q_low = OnnxGather(quantiles_init, current_index, axis=0, op_version=opv)
+        q_high = OnnxGather(quantiles_init, index_plus_one, axis=0, op_version=opv)
+        r_low = OnnxGather(references_init, current_index, axis=0, op_version=opv)
+        r_high = OnnxGather(references_init, index_plus_one, axis=0, op_version=opv)
+
+        # Handle boundary conditions
+        mask_low = OnnxLessOrEqual(flat_input, q_low, op_version=opv)
+        mask_high = OnnxGreaterOrEqual(flat_input, q_high, op_version=opv)
+
+        delta_q = OnnxSub(q_high, q_low, op_version=opv)
+        delta_r = OnnxSub(r_high, r_low, op_version=opv)
+
+        slope = OnnxDiv(delta_r, delta_q, op_version=opv)
+        invalid_slope = OnnxOr(
+            OnnxLessOrEqual(delta_q, np.array(0, dtype=dtype), op_version=opv),
+            OnnxIsNaN(slope, op_version=opv),
+            op_version=opv,
+        )
+        slope = OnnxWhere(
+            invalid_slope, np.array(0, dtype=dtype), slope, op_version=opv
+        )
+
+        interp_normal = OnnxAdd(
+            r_low,
+            OnnxMul(slope, OnnxSub(flat_input, q_low, op_version=opv), op_version=opv),
+            op_version=opv,
+        )
+
+        interp = OnnxWhere(
+            mask_high,
+            r_high,
+            OnnxWhere(mask_low, r_low, interp_normal, op_version=opv),
+            op_version=opv,
+        )
+        return interp
+
+    # First interpolation
+    interp1 = create_interpolation(X_clean, quantiles_name, references_name)
+
+    # Second interpolation
+    X_neg = OnnxNeg(X_clean, op_version=opv)
+    X_neg.add_to(scope, container)
+    interp2 = create_interpolation(
+        X_neg, quantiles_neg_rev_name, references_neg_rev_name
+    )
+
+    # Combine results
+    final_interp = OnnxMul(
+        OnnxSub(interp1, interp2, op_version=opv),
+        np.array(0.5, dtype=dtype),
+        op_version=opv,
+    )
+
+    # Reshape to original shape
+    output_shape = OnnxShape(X_clean, op_version=opv)
+    final = OnnxReshape(
+        final_interp, output_shape, op_version=opv, output_names=[op_out]
+    )
+    min_quantiles = quantiles[0, :]  # Minimum values of quantiles
+    max_quantiles = quantiles[-1, :]  # Maximum values of quantiles
+
+    # Comparison with minimum and maximum values
+    min_mask = OnnxFlatten(
+        OnnxEqual(X_clean, min_quantiles, op_version=opv), axis=0, op_version=opv
+    )
+    max_mask = OnnxFlatten(
+        OnnxEqual(X_clean, max_quantiles, op_version=opv), axis=0, op_version=opv
+    )
+
+    # Set 0 for minimum values and 1 for maximum values
+    final = OnnxWhere(
+        max_mask,
+        np.array(1.0, dtype=dtype),
+        OnnxWhere(
+            min_mask,
+            OnnxMul(final_interp, np.array(0.0, dtype=dtype), op_version=opv),
+            final_interp,
+            op_version=opv,
+        ),
+        op_version=opv,
+    )
+    final = OnnxReshape(final, output_shape, op_version=opv, output_names=[op_out])
+    final.add_to(scope, container)
+
+
+register_converter("SklearnQuantileTransformer", convert_quantile_transformer)
diff --git a/skl2onnx/shape_calculators/__init__.py b/skl2onnx/shape_calculators/__init__.py
index 3cea54b..fdb3acb 100644
--- a/skl2onnx/shape_calculators/__init__.py
+++ b/skl2onnx/shape_calculators/__init__.py
@@ -40,6 +40,7 @@ from . import pipelines
 from . import polynomial_features
 from . import power_transformer
 from . import quadratic_discriminant_analysis
+from . import quantile_transformer
 from . import random_projection
 from . import random_trees_embedding
 from . import replace_op
@@ -94,6 +95,7 @@ __all__ = [
     polynomial_features,
     power_transformer,
     quadratic_discriminant_analysis,
+    quantile_transformer,
     random_projection,
     random_trees_embedding,
     replace_op,
diff --git a/skl2onnx/shape_calculators/quantile_transformer.py b/skl2onnx/shape_calculators/quantile_transformer.py
new file mode 100644
index 0000000..c77aaee
--- /dev/null
+++ b/skl2onnx/shape_calculators/quantile_transformer.py
@@ -0,0 +1,24 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import copy
+from ..common._registration import register_shape_calculator
+from ..common.utils import check_input_and_output_numbers, check_input_and_output_types
+from ..common.data_types import FloatTensorType, Int64TensorType, DoubleTensorType
+
+
+def quantile_transformer_shape_calculator(operator):
+    """Shape calculator for QuantileTransformer"""
+    check_input_and_output_numbers(operator, output_count_range=1)
+    check_input_and_output_types(
+        operator, good_input_types=[FloatTensorType, Int64TensorType, DoubleTensorType]
+    )
+
+    N = operator.inputs[0].get_first_dimension()
+    model = operator.raw_operator
+    operator.outputs[0].type = copy.deepcopy(operator.inputs[0].type)
+    operator.outputs[0].type.shape = [N, model.quantiles_.shape[1]]
+
+
+register_shape_calculator(
+    "SklearnQuantileTransformer", quantile_transformer_shape_calculator
+)


Test patch:
diff --git a/tests/test_sklearn_quantile_transformer.py b/tests/test_sklearn_quantile_transformer.py
new file mode 100644
index 0000000..d7b8617
--- /dev/null
+++ b/tests/test_sklearn_quantile_transformer.py
@@ -0,0 +1,64 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Tests scikit-learn's polynomial features converter.
+"""
+import unittest
+import numpy as np
+from sklearn.preprocessing import QuantileTransformer
+from skl2onnx import convert_sklearn
+from skl2onnx.common.data_types import FloatTensorType
+from test_utils import dump_data_and_model
+
+
+class TestSklearnQuantileTransformer(unittest.TestCase):
+    def test_quantile_transformer_simple(self):
+        X = np.empty((100, 2), dtype=np.float32)
+        X[:, 0] = np.arange(X.shape[0])
+        X[:, 1] = np.arange(X.shape[0]) * 2
+        model = QuantileTransformer(n_quantiles=6).fit(X)
+        model_onnx = convert_sklearn(
+            model, "test", [("input", FloatTensorType([None, X.shape[1]]))]
+        )
+        self.assertTrue(model_onnx is not None)
+        dump_data_and_model(
+            X.astype(np.float32),
+            model,
+            model_onnx,
+            basename="SklearnQuantileTransformer",
+        )
+
+    def test_quantile_transformer_int(self):
+        X = np.random.randint(0, 5, (100, 20))
+        model = QuantileTransformer(n_quantiles=6).fit(X)
+        model_onnx = convert_sklearn(
+            model, "test", [("input", FloatTensorType([None, X.shape[1]]))]
+        )
+        self.assertTrue(model_onnx is not None)
+        dump_data_and_model(
+            X.astype(np.float32),
+            model,
+            model_onnx,
+            basename="SklearnQuantileTransformer",
+        )
+
+    def test_quantile_transformer_nan(self):
+        X = np.random.randint(0, 5, (100, 20))
+        X = X.astype(np.float32)
+        X[0][0] = np.nan
+        X[1][1] = np.nan
+        model = QuantileTransformer(n_quantiles=6).fit(X)
+        model_onnx = convert_sklearn(
+            model, "test", [("input", FloatTensorType([None, X.shape[1]]))]
+        )
+        self.assertTrue(model_onnx is not None)
+        dump_data_and_model(
+            X.astype(np.float32),
+            model,
+            model_onnx,
+            basename="SklearnQuantileTransformer",
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()