onnx--onnx/onnx/test/reference_evaluator_ml_test.py

# Copyright (c) ONNX Project Contributors

# SPDX-License-Identifier: Apache-2.0
# mypy: ignore-errors
from __future__ import annotations

import importlib
import itertools
import os

import numpy as np
import pytest
from numpy.testing import assert_allclose

import onnx
from onnx import ONNX_ML, TensorProto, TypeProto, ValueInfoProto
from onnx.helper import (
    make_graph,
    make_model_gen_version,
    make_node,
    make_opsetid,
    make_tensor,
    make_tensor_value_info,
)
from onnx.reference import ReferenceEvaluator
from onnx.reference.ops.aionnxml.op_tree_ensemble import (
    AggregationFunction,
    Mode,
    PostTransform,
)

# TODO (https://github.com/microsoft/onnxruntime/issues/14932): Get max supported version from onnxruntime directly
# For now, bump the version in CIs whenever there is a new onnxruntime release
ORT_MAX_IR_SUPPORTED_VERSION = int(os.getenv("ORT_MAX_IR_SUPPORTED_VERSION", "8"))
ORT_MAX_ONNX_OPSET_SUPPORTED_VERSION = int(
    os.getenv("ORT_MAX_ONNX_OPSET_SUPPORTED_VERSION", "18")
)
ORT_MAX_ML_OPSET_SUPPORTED_VERSION = int(
    os.getenv("ORT_MAX_ML_OPSET_SUPPORTED_VERSION", "3")
)

TARGET_OPSET = onnx.defs.onnx_opset_version() - 2
TARGET_OPSET_ML = onnx.defs.onnx_ml_opset_version()
OPSETS = [make_opsetid("", TARGET_OPSET), make_opsetid("ai.onnx.ml", TARGET_OPSET_ML)]


def has_onnxruntime():
    return importlib.util.find_spec("onnxruntime") is not None


class TestReferenceEvaluatorAiOnnxMl:
    @staticmethod
    def _check_ort(model, feeds, atol=0, rtol=0, equal=False, rev=False):
        if not has_onnxruntime():
            pytest.skip("onnxruntime not installed")

        from onnxruntime import InferenceSession  # noqa: PLC0415

        onnx_domain_opset = ORT_MAX_ONNX_OPSET_SUPPORTED_VERSION
        ml_domain_opset = ORT_MAX_ML_OPSET_SUPPORTED_VERSION
        for opset in model.opset_import:
            if opset.domain in ("", "ai.onnx"):
                onnx_domain_opset = opset.version
                break
        for opset in model.opset_import:
            if opset.domain == "ai.onnx.ml":
                ml_domain_opset = opset.version
                break
        # The new IR or opset version is not supported by onnxruntime yet
        if (
            model.ir_version > ORT_MAX_IR_SUPPORTED_VERSION
            or onnx_domain_opset > ORT_MAX_ONNX_OPSET_SUPPORTED_VERSION
            or ml_domain_opset > ORT_MAX_ML_OPSET_SUPPORTED_VERSION
        ):
            return

        ort = InferenceSession(
            model.SerializeToString(), providers=["CPUExecutionProvider"]
        )
        sess = ReferenceEvaluator(model)
        expected = ort.run(None, feeds)
        got = sess.run(None, feeds)
        if len(expected) != len(got):
            raise AssertionError(
                f"onnxruntime returns a different number of output "
                f"{len(expected)} != {len(sess)} (ReferenceEvaluator)."
            )
        look = (
            zip(reversed(expected), reversed(got), strict=True)
            if rev
            else zip(expected, got, strict=True)
        )
        for i, (e, g) in enumerate(look):
            if e.shape != g.shape:
                raise AssertionError(
                    f"Unexpected shape {g.shape} for output {i} "
                    f"(expecting {e.shape})\n{e!r}\n---\n{g!r}."
                )
            if equal:
                if e.tolist() != g.tolist():
                    raise AssertionError(
                        f"Discrepancies for output {i}"
                        f"\nexpected=\n{e}\n!=\nresults=\n{g}"
                    )
            else:
                assert_allclose(
                    actual=g,
                    desired=e,
                    atol=atol,
                    rtol=rtol,
                    err_msg=f"Discrepancies for output {i} expected[0]={e.ravel()[0]}.",
                )

    @pytest.mark.skipif(not ONNX_ML, reason="onnx not compiled with ai.onnx.ml")
    def test_binarizer(self):
        X = make_tensor_value_info("X", TensorProto.FLOAT, [None, None])
        Y = make_tensor_value_info("Y", TensorProto.FLOAT, [None, None])
        node1 = make_node("Binarizer", ["X"], ["Y"], threshold=5.5, domain="ai.onnx.ml")
        graph = make_graph([node1], "ml", [X], [Y])
        model = make_model_gen_version(graph, opset_imports=OPSETS)
        onnx.checker.check_model(model)
        x = np.arange(12).reshape((3, 4)).astype(np.float32)
        expected = np.array(
            [[0, 0, 0, 0], [0, 0, 1, 1], [1, 1, 1, 1]], dtype=np.float32
        )
        self._check_ort(model, {"X": x})
        sess = ReferenceEvaluator(model)
        got = sess.run(None, {"X": x})[0]
        assert_allclose(got, expected)

    @pytest.mark.skipif(not ONNX_ML, reason="onnx not compiled with ai.onnx.ml")
    def test_scaler(self):
        X = make_tensor_value_info("X", TensorProto.FLOAT, [None, None])
        Y = make_tensor_value_info("Y", TensorProto.FLOAT, [None, None])
        node1 = make_node(
            "Scaler", ["X"], ["Y"], scale=[0.5], offset=[-4.5], domain="ai.onnx.ml"
        )
        graph = make_graph([node1], "ml", [X], [Y])
        model = make_model_gen_version(graph, opset_imports=OPSETS)
        onnx.checker.check_model(model)
        x = np.arange(12).reshape((3, 4)).astype(np.float32)
        expected = np.array(
            [
                [2.25, 2.75, 3.25, 3.75],
                [4.25, 4.75, 5.25, 5.75],
                [6.25, 6.75, 7.25, 7.75],
            ],
            dtype=np.float32,
        )
        self._check_ort(model, {"X": x})
        sess = ReferenceEvaluator(model)
        got = sess.run(None, {"X": x})[0]
        assert_allclose(got, expected)

    @pytest.mark.skipif(not ONNX_ML, reason="onnx not compiled with ai.onnx.ml")
    def test_array_feature_extractor(self):
        X = make_tensor_value_info("X", TensorProto.FLOAT, [None, None])
        A = make_tensor_value_info("A", TensorProto.INT64, [None])
        Y = make_tensor_value_info("Y", TensorProto.FLOAT, [None, None])
        node1 = make_node(
            "ArrayFeatureExtractor", ["X", "A"], ["Y"], domain="ai.onnx.ml"
        )
        graph = make_graph([node1], "ml", [X, A], [Y])
        model = make_model_gen_version(graph, opset_imports=OPSETS)
        onnx.checker.check_model(model)
        x = np.arange(12).reshape((3, 4)).astype(np.float32)

        expected = np.array([[0, 4, 8]], dtype=np.float32).T
        feeds = {"X": x, "A": np.array([0], dtype=np.int64)}
        self._check_ort(model, feeds)
        sess = ReferenceEvaluator(model)
        got = sess.run(None, feeds)[0]
        assert_allclose(got, expected)

        expected = np.array([[0, 4, 8], [1, 5, 9]], dtype=np.float32).T
        feeds = {"X": x, "A": np.array([0, 1], dtype=np.int64)}
        self._check_ort(model, feeds)
        sess = ReferenceEvaluator(model)
        got = sess.run(None, feeds)[0]
        assert_allclose(got, expected)

        expected = np.array(
            [[0, 4, 8], [1, 5, 9], [0, 4, 8], [1, 5, 9], [0, 4, 8], [1, 5, 9]],
            dtype=np.float32,
        ).T
        feeds = {"X": x, "A": np.array([0, 1, 0, 1, 0, 1], dtype=np.int64)}
        self._check_ort(model, feeds)
        sess = ReferenceEvaluator(model)
        got = sess.run(None, feeds)[0]
        assert_allclose(got, expected)

    @pytest.mark.parametrize(
        "norm, compute",
        [
            ("MAX", lambda x: x / x.max(axis=1, keepdims=1)),
            ("L1", lambda x: x / np.abs(x).sum(axis=1, keepdims=1)),
            ("L2", lambda x: x / (x**2).sum(axis=1, keepdims=1) ** 0.5),
        ],
    )
    @pytest.mark.skipif(not ONNX_ML, reason="onnx not compiled with ai.onnx.ml")
    def test_normalizer(self, norm, compute):
        X = make_tensor_value_info("X", TensorProto.FLOAT, [None, None])
        Y = make_tensor_value_info("Y", TensorProto.FLOAT, [None, None])
        x = np.arange(12).reshape((3, 4)).astype(np.float32)

        node1 = make_node("Normalizer", ["X"], ["Y"], norm=norm, domain="ai.onnx.ml")
        graph = make_graph([node1], "ml", [X], [Y])
        model = make_model_gen_version(graph, opset_imports=OPSETS)
        onnx.checker.check_model(model)

        feeds = {"X": x}
        expected = compute(x)
        self._check_ort(model, feeds, atol=1e-6)
        sess = ReferenceEvaluator(model)
        got = sess.run(None, feeds)[0]
        assert_allclose(got, expected, atol=1e-6)

    @pytest.mark.parametrize(
        "inputdimensions, expected_value",
        [
            ((1,), np.array([[0], [3], [6]], dtype=np.float32)),
            ((2,), np.array([[0, 1], [3, 4], [6, 7]], dtype=np.float32)),
            (
                (4,),
                np.array([[0, 1, 2, 0], [3, 4, 5, 0], [6, 7, 8, 0]], dtype=np.float32),
            ),
            ((1, 1), np.array([[0, 0.5], [3, 3.5], [6, 6.5]], dtype=np.float32)),
            ((0, 1), np.array([[0.5], [3.5], [6.5]], dtype=np.float32)),
        ],
    )
    @pytest.mark.skipif(not ONNX_ML, reason="onnx not compiled with ai.onnx.ml")
    def test_feature_vectorizer(self, inputdimensions, expected_value):
        X = [
            make_tensor_value_info("X0", TensorProto.FLOAT, [None, None]),
            make_tensor_value_info("X1", TensorProto.FLOAT, [None, None]),
        ]
        Y = make_tensor_value_info("Y", TensorProto.FLOAT, [None, None])
        x = [
            np.arange(9).reshape((3, 3)).astype(np.float32),
            np.arange(9).reshape((3, 3)).astype(np.float32) + 0.5,
        ]
        att = (
            list(inputdimensions)
            if isinstance(inputdimensions, tuple)
            else inputdimensions
        )
        node1 = make_node(
            "FeatureVectorizer",
            [f"X{i}" for i in range(len(att))],
            ["Y"],
            inputdimensions=att,
            domain="ai.onnx.ml",
        )
        graph = make_graph([node1], "ml", X[: len(att)], [Y])
        model = make_model_gen_version(graph, opset_imports=OPSETS)
        onnx.checker.check_model(model)

        feeds = {f"X{i}": v for i, v in enumerate(x[: len(att)])}
        self._check_ort(model, feeds, atol=1e-6)
        sess = ReferenceEvaluator(model)
        got = sess.run(None, feeds)[0]
        assert_allclose(got, expected_value, atol=1e-6)

    @pytest.mark.skipif(not ONNX_ML, reason="onnx not compiled with ai.onnx.ml")
    def test_imputer_float(self):
        X = make_tensor_value_info("X", TensorProto.FLOAT, [None, None])
        Y = make_tensor_value_info("Y", TensorProto.FLOAT, [None, None])
        node1 = make_node(
            "Imputer",
            ["X"],
            ["Y"],
            domain="ai.onnx.ml",
            imputed_value_floats=np.array([0], dtype=np.float32),
            replaced_value_float=np.nan,
        )
        graph = make_graph([node1], "ml", [X], [Y])
        model = make_model_gen_version(graph, opset_imports=OPSETS)
        onnx.checker.check_model(model)
        x = np.array([[0, 1, np.nan, 3]], dtype=np.float32).T
        expected = np.array([[0, 1, 0, 3]], dtype=np.float32).T
        self._check_ort(model, {"X": x})
        sess = ReferenceEvaluator(model)
        got = sess.run(None, {"X": x})[0]
        assert_allclose(got, expected)

    @pytest.mark.skipif(not ONNX_ML, reason="onnx not compiled with ai.onnx.ml")
    def test_imputer_float_2d(self):
        X = make_tensor_value_info("X", TensorProto.FLOAT, [None, None])
        Y = make_tensor_value_info("Y", TensorProto.FLOAT, [None, None])
        node1 = make_node(
            "Imputer",
            ["X"],
            ["Y"],
            domain="ai.onnx.ml",
            imputed_value_floats=np.array([0, 0.1], dtype=np.float32),
            replaced_value_float=np.nan,
        )
        graph = make_graph([node1], "ml", [X], [Y])
        model = make_model_gen_version(graph, opset_imports=OPSETS)
        onnx.checker.check_model(model)
        x = np.array([[0, 1, np.nan, 3], [0, 1, np.nan, 3]], dtype=np.float32).T
        expected = np.array([[0, 1, 0, 3], [0, 1, 0.1, 3]], dtype=np.float32).T
        self._check_ort(model, {"X": x})
        sess = ReferenceEvaluator(model)
        got = sess.run(None, {"X": x})[0]
        assert_allclose(got, expected)

    @pytest.mark.skipif(not ONNX_ML, reason="onnx not compiled with ai.onnx.ml")
    def test_imputer_int(self):
        X = make_tensor_value_info("X", TensorProto.INT64, [None, None])
        Y = make_tensor_value_info("Y", TensorProto.INT64, [None, None])
        node1 = make_node(
            "Imputer",
            ["X"],
            ["Y"],
            domain="ai.onnx.ml",
            imputed_value_int64s=np.array([0], dtype=np.int64),
            replaced_value_int64=-1,
        )
        graph = make_graph([node1], "ml", [X], [Y])
        model = make_model_gen_version(graph, opset_imports=OPSETS)
        onnx.checker.check_model(model)
        x = np.array([[0, 1, -1, 3]], dtype=np.int64).T
        expected = np.array([[0, 1, 0, 3]], dtype=np.int64).T
        self._check_ort(model, {"X": x})
        sess = ReferenceEvaluator(model)
        got = sess.run(None, {"X": x})[0]
        assert_allclose(got, expected)

    @pytest.mark.skipif(not ONNX_ML, reason="onnx not compiled with ai.onnx.ml")
    def test_label_encoder_float_int(self):
        X = make_tensor_value_info("X", TensorProto.FLOAT, [None, None])
        Y = make_tensor_value_info("Y", TensorProto.INT64, [None, None])
        node1 = make_node(
            "LabelEncoder",
            ["X"],
            ["Y"],
            domain="ai.onnx.ml",
            default_int64=-5,
            keys_floats=[4.0, 1.0, 2.0, 3.0],
            values_int64s=[0, 1, 2, 3],
        )
        graph = make_graph([node1], "ml", [X], [Y])
        model = make_model_gen_version(graph, opset_imports=OPSETS)
        onnx.checker.check_model(model)
        x = np.array([[0, 1, np.nan, 3, 4]], dtype=np.float32).T
        expected = np.array([[-5, 1, -5, 3, 0]], dtype=np.int64).T
        self._check_ort(model, {"X": x})
        sess = ReferenceEvaluator(model)
        got = sess.run(None, {"X": x})[0]
        assert_allclose(got, expected)

    @pytest.mark.skipif(not ONNX_ML, reason="onnx not compiled with ai.onnx.ml")
    def test_label_encoder_int_string(self):
        X = make_tensor_value_info("X", TensorProto.INT64, [None, None])
        Y = make_tensor_value_info("Y", TensorProto.STRING, [None, None])
        node1 = make_node(
            "LabelEncoder",
            ["X"],
            ["Y"],
            domain="ai.onnx.ml",
            default_string="NONE",
            keys_int64s=[1, 2, 3, 4],
            values_strings=["a", "b", "cc", "ddd"],
        )
        graph = make_graph([node1], "ml", [X], [Y])
        model = make_model_gen_version(graph, opset_imports=OPSETS)
        onnx.checker.check_model(model)
        x = np.array([[0, 1, 3, 4]], dtype=np.int64).T
        expected = np.array([["NONE"], ["a"], ["cc"], ["ddd"]])
        self._check_ort(model, {"X": x}, equal=True)
        sess = ReferenceEvaluator(model)
        got = sess.run(None, {"X": x})[0]
        assert expected.tolist() == got.tolist()

    @pytest.mark.skipif(not ONNX_ML, reason="onnx not compiled with ai.onnx.ml")
    def test_label_encoder_int_string_tensor_attributes(self):
        X = make_tensor_value_info("X", TensorProto.INT64, [None, None])
        Y = make_tensor_value_info("Y", TensorProto.STRING, [None, None])
        node = make_node(
            "LabelEncoder",
            ["X"],
            ["Y"],
            domain="ai.onnx.ml",
            keys_tensor=make_tensor(
                "keys_tensor", TensorProto.INT64, [4], [1, 2, 3, 4]
            ),
            values_tensor=make_tensor(
                "values_tensor", TensorProto.STRING, [4], ["a", "b", "cc", "ddd"]
            ),
            default_tensor=make_tensor(
                "default_tensor", TensorProto.STRING, [], ["NONE"]
            ),
        )
        graph = make_graph([node], "ml", [X], [Y])
        model = make_model_gen_version(graph, opset_imports=OPSETS)
        onnx.checker.check_model(model)
        x = np.array([[0, 1, 3, 4]], dtype=np.int64).T
        expected = np.array([["NONE"], ["a"], ["cc"], ["ddd"]])
        self._check_ort(model, {"X": x}, equal=True)
        sess = ReferenceEvaluator(model)
        got = sess.run(None, {"X": x})[0]
        assert expected.tolist() == got.tolist()

    @pytest.mark.skipif(not ONNX_ML, reason="onnx not compiled with ai.onnx.ml")
    def test_dict_vectorizer(self):
        value_type = TypeProto()
        value_type.tensor_type.elem_type = TensorProto.INT64
        onnx_type = TypeProto()
        onnx_type.map_type.key_type = TensorProto.STRING
        onnx_type.map_type.value_type.CopyFrom(value_type)
        value_info = ValueInfoProto()
        value_info.name = "X"
        value_info.type.CopyFrom(onnx_type)

        X = value_info
        Y = make_tensor_value_info("Y", TensorProto.INT64, [None, None])
        node1 = make_node(
            "DictVectorizer",
            ["X"],
            ["Y"],
            domain="ai.onnx.ml",
            string_vocabulary=["a", "c", "b", "z"],
        )
        graph = make_graph([node1], "ml", [X], [Y])
        model = make_model_gen_version(graph, opset_imports=OPSETS)
        onnx.checker.check_model(model)
        x = {"a": np.array(4, dtype=np.int64), "c": np.array(8, dtype=np.int64)}
        expected = np.array([4, 8, 0, 0], dtype=np.int64)
        # Unexpected input data type. Actual: ((map(string,tensor(float)))) , expected: ((map(string,tensor(int64))))
        # self._check_ort(model, {"X": x}, equal=True)
        sess = ReferenceEvaluator(model)
        got = sess.run(None, {"X": x})[0]
        assert expected.tolist() == got.tolist()

    @pytest.mark.skipif(not ONNX_ML, reason="onnx not compiled with ai.onnx.ml")
    def test_one_hot_encoder_int(self):
        X = make_tensor_value_info("X", TensorProto.INT64, [None, None])
        Y = make_tensor_value_info("Y", TensorProto.FLOAT, [None, None, None])
        node1 = make_node(
            "OneHotEncoder",
            ["X"],
            ["Y"],
            domain="ai.onnx.ml",
            zeros=1,
            cats_int64s=[1, 2, 3],
        )
        graph = make_graph([node1], "ml", [X], [Y])
        model = make_model_gen_version(graph, opset_imports=OPSETS)
        onnx.checker.check_model(model)
        x = np.array([[5, 1, 3], [2, 1, 3]], dtype=np.int64)
        expected = np.array(
            [[[0, 0, 0], [1, 0, 0], [0, 0, 1]], [[0, 1, 0], [1, 0, 0], [0, 0, 1]]],
            dtype=np.float32,
        )
        self._check_ort(model, {"X": x}, equal=True)
        sess = ReferenceEvaluator(model)
        got = sess.run(None, {"X": x})[0]
        assert expected.tolist() == got.tolist()

    @pytest.mark.skipif(not ONNX_ML, reason="onnx not compiled with ai.onnx.ml")
    def test_one_hot_encoder_string(self):
        X = make_tensor_value_info("X", TensorProto.STRING, [None, None])
        Y = make_tensor_value_info("Y", TensorProto.FLOAT, [None, None, None])
        node1 = make_node(
            "OneHotEncoder",
            ["X"],
            ["Y"],
            domain="ai.onnx.ml",
            zeros=1,
            cats_strings=["c1", "c2", "c3"],
        )
        graph = make_graph([node1], "ml", [X], [Y])
        model = make_model_gen_version(graph, opset_imports=OPSETS)
        onnx.checker.check_model(model)
        x = np.array([["c5", "c1", "c3"], ["c2", "c1", "c3"]])
        expected = np.array(
            [[[0, 0, 0], [1, 0, 0], [0, 0, 1]], [[0, 1, 0], [1, 0, 0], [0, 0, 1]]],
            dtype=np.float32,
        )
        self._check_ort(model, {"X": x}, equal=True)
        sess = ReferenceEvaluator(model)
        got = sess.run(None, {"X": x})[0]
        assert expected.tolist() == got.tolist()

    @pytest.mark.skipif(not ONNX_ML, reason="onnx not compiled with ai.onnx.ml")
    def test_one_hot_encoder_zeros(self):
        X = make_tensor_value_info("X", TensorProto.INT64, [None, None])
        Y = make_tensor_value_info("Y", TensorProto.FLOAT, [None, None, None])
        node1 = make_node(
            "OneHotEncoder",
            ["X"],
            ["Y"],
            domain="ai.onnx.ml",
            zeros=0,
            cats_int64s=[1, 2, 3],
        )
        graph = make_graph([node1], "ml", [X], [Y])
        model = make_model_gen_version(graph, opset_imports=OPSETS)
        onnx.checker.check_model(model)
        x = np.array([[2, 1, 3], [2, 1, 3]], dtype=np.int64)
        expected = np.array(
            [[[0, 1, 0], [1, 0, 0], [0, 0, 1]], [[0, 1, 0], [1, 0, 0], [0, 0, 1]]],
            dtype=np.float32,
        )
        self._check_ort(model, {"X": x}, equal=True)
        sess = ReferenceEvaluator(model)
        got = sess.run(None, {"X": x})[0]
        assert expected.tolist() == got.tolist()

    @pytest.mark.skipif(not ONNX_ML, reason="onnx not compiled with ai.onnx.ml")
    def test_linear_regressor(self):
        X = make_tensor_value_info("X", TensorProto.FLOAT, [None, None])
        Y = make_tensor_value_info("Y", TensorProto.FLOAT, [None, None])
        node1 = make_node(
            "LinearRegressor",
            ["X"],
            ["Y"],
            domain="ai.onnx.ml",
            coefficients=[0.3, -0.77],
            intercepts=[0.5],
            post_transform="NONE",
            targets=1,
        )
        graph = make_graph([node1], "ml", [X], [Y])
        model = make_model_gen_version(graph, opset_imports=OPSETS)
        onnx.checker.check_model(model)
        x = np.arange(6).reshape((-1, 2)).astype(np.float32)
        expected = np.array([[-0.27], [-1.21], [-2.15]], dtype=np.float32)
        self._check_ort(model, {"X": x}, equal=True)
        sess = ReferenceEvaluator(model)
        got = sess.run(None, {"X": x})
        assert_allclose(got[0], expected, atol=1e-6)

    @pytest.mark.skipif(not ONNX_ML, reason="onnx not compiled with ai.onnx.ml")
    def test_linear_regressor_2(self):
        X = make_tensor_value_info("X", TensorProto.FLOAT, [None, None])
        Y = make_tensor_value_info("Y", TensorProto.FLOAT, [None, None])
        node1 = make_node(
            "LinearRegressor",
            ["X"],
            ["Y"],
            domain="ai.onnx.ml",
            coefficients=[0.3, -0.77, 0.3, -0.77],
            intercepts=[0.5, 0.7],
            post_transform="NONE",
            targets=2,
        )
        graph = make_graph([node1], "ml", [X], [Y])
        model = make_model_gen_version(graph, opset_imports=OPSETS)
        onnx.checker.check_model(model)
        x = np.arange(6).reshape((-1, 2)).astype(np.float32)
        expected = np.array(
            [[-0.27, -0.07], [-1.21, -1.01], [-2.15, -1.95]], dtype=np.float32
        )
        self._check_ort(model, {"X": x}, equal=True)
        sess = ReferenceEvaluator(model)
        got = sess.run(None, {"X": x})
        assert_allclose(got[0], expected, atol=1e-6)

    @pytest.mark.parametrize(
        "post", ["SOFTMAX", "NONE", "LOGISTIC", "SOFTMAX_ZERO", "PROBIT"]
    )
    @pytest.mark.skipif(not ONNX_ML, reason="onnx not compiled with ai.onnx.ml")
    def test_linear_classifier_multi(self, post):
        X = make_tensor_value_info("X", TensorProto.FLOAT, [None, None])
        In = make_tensor_value_info("I", TensorProto.INT64, [None])
        Y = make_tensor_value_info("Y", TensorProto.FLOAT, [None, None])
        expected_post = {
            "NONE": [
                np.array([0, 2, 2], dtype=np.int64),
                np.array(
                    [[2.41, -2.12, 0.59], [0.67, -1.14, 1.35], [-1.07, -0.16, 2.11]],
                    dtype=np.float32,
                ),
            ],
            "LOGISTIC": [
                np.array([0, 2, 2], dtype=np.int64),
                np.array(
                    [
                        [0.917587, 0.107168, 0.643365],
                        [0.661503, 0.24232, 0.79413],
                        [0.255403, 0.460085, 0.891871],
                    ],
                    dtype=np.float32,
                ),
            ],
            "SOFTMAX": [
                np.array([0, 2, 2], dtype=np.int64),
                np.array(
                    [
                        [0.852656, 0.009192, 0.138152],
                        [0.318722, 0.05216, 0.629118],
                        [0.036323, 0.090237, 0.87344],
                    ],
                    dtype=np.float32,
                ),
            ],
            "SOFTMAX_ZERO": [
                np.array([0, 2, 2], dtype=np.int64),
                np.array(
                    [
                        [0.852656, 0.009192, 0.138152],
                        [0.318722, 0.05216, 0.629118],
                        [0.036323, 0.090237, 0.87344],
                    ],
                    dtype=np.float32,
                ),
            ],
            "PROBIT": [
                np.array([1, 1, 1], dtype=np.int64),
                np.array(
                    [
                        [-0.527324, -0.445471, -1.080504],
                        [-0.067731, 0.316014, -0.310748],
                        [0.377252, 1.405167, 0.295001],
                    ],
                    dtype=np.float32,
                ),
            ],
        }
        if post == "PROBIT":
            coefficients = [0.058, 0.029, 0.09, 0.058, 0.029, 0.09]
            intercepts = [0.27, 0.27, 0.05]
        else:
            coefficients = [-0.58, -0.29, -0.09, 0.58, 0.29, 0.09]
            intercepts = [2.7, -2.7, 0.5]
        node1 = make_node(
            "LinearClassifier",
            ["X"],
            ["I", "Y"],
            domain="ai.onnx.ml",
            classlabels_ints=[0, 1, 2],
            coefficients=coefficients,
            intercepts=intercepts,
            multi_class=0,
            post_transform=post,
        )
        graph = make_graph([node1], "ml", [X], [In, Y])
        model = make_model_gen_version(graph, opset_imports=OPSETS)
        onnx.checker.check_model(model)
        x = np.arange(6).reshape((-1, 2)).astype(np.float32)
        self._check_ort(model, {"X": x}, rev=True, atol=1e-4)
        sess = ReferenceEvaluator(model)
        got = sess.run(None, {"X": x})
        expected = expected_post[post]
        assert_allclose(got[1], expected[1], atol=1e-4)
        assert_allclose(got[0], expected[0])

    @pytest.mark.parametrize("post", ["SOFTMAX", "NONE", "LOGISTIC", "SOFTMAX_ZERO"])
    @pytest.mark.skipif(not ONNX_ML, reason="onnx not compiled with ai.onnx.ml")
    def test_linear_classifier_binary(self, post):
        X = make_tensor_value_info("X", TensorProto.FLOAT, [None, None])
        In = make_tensor_value_info("I", TensorProto.INT64, [None])
        Y = make_tensor_value_info("Y", TensorProto.FLOAT, [None, None])
        expected_post = {
            "NONE": [
                np.array([1, 1], dtype=np.int64),
                np.array([[-9.53, 9.53], [-6.65, 6.65]], dtype=np.float32),
            ],
            "LOGISTIC": [
                np.array([1, 1], dtype=np.int64),
                np.array(
                    [[7.263436e-05, 9.999274e-01], [1.292350e-03, 9.987077e-01]],
                    dtype=np.float32,
                ),
            ],
            "SOFTMAX": [
                np.array([1, 1], dtype=np.int64),
                np.array(
                    [[5.276517e-09, 1.000000e00], [1.674492e-06, 9.999983e-01]],
                    dtype=np.float32,
                ),
            ],
            "SOFTMAX_ZERO": [
                np.array([1, 1], dtype=np.int64),
                np.array(
                    [[5.276517e-09, 1.000000e00], [1.674492e-06, 9.999983e-01]],
                    dtype=np.float32,
                ),
            ],
        }
        x = np.arange(6).reshape((-1, 3)).astype(np.float32)
        expected = expected_post[post]
        node1 = make_node(
            "LinearClassifier",
            ["X"],
            ["I", "Y"],
            domain="ai.onnx.ml",
            classlabels_ints=[0, 1],
            coefficients=[-0.58, -0.29, -0.09],
            intercepts=[10.0],
            multi_class=0,
            post_transform=post,
        )
        graph = make_graph([node1], "ml", [X], [In, Y])
        model = make_model_gen_version(graph, opset_imports=OPSETS)
        onnx.checker.check_model(model)
        # onnxruntime answer seems odd.
        # self._check_ort(model, {"X": x}, rev=True)
        sess = ReferenceEvaluator(model)
        got = sess.run(None, {"X": x})
        assert_allclose(got[1], expected[1], atol=1e-6)
        assert_allclose(got[0], expected[0])

    @pytest.mark.skipif(not ONNX_ML, reason="onnx not compiled with ai.onnx.ml")
    @pytest.mark.parametrize(
        "post, expected",
        [
            (
                "NONE",
                (
                    np.array([1, 0], dtype=np.int64),
                    np.array([[2.23], [-0.65]], dtype=np.float32),
                ),
            ),
            (
                "LOGISTIC",
                (
                    np.array([1, 0], dtype=np.int64),
                    np.array([[0.902911], [0.34299]], dtype=np.float32),
                ),
            ),
            (
                "SOFTMAX",
                (
                    np.array([1, 1], dtype=np.int64),
                    np.array([[1.0], [1.0]], dtype=np.float32),
                ),
            ),
            (
                "SOFTMAX_ZERO",
                (
                    np.array([1, 1], dtype=np.int64),
                    np.array([[1.0], [1.0]], dtype=np.float32),
                ),
            ),
        ],
    )
    def test_linear_classifier_unary(self, post, expected):
        X = make_tensor_value_info("X", TensorProto.FLOAT, [None, None])
        In = make_tensor_value_info("I", TensorProto.INT64, [None])
        Y = make_tensor_value_info("Y", TensorProto.FLOAT, [None, None])
        x = np.arange(6).reshape((-1, 3)).astype(np.float32)
        node1 = make_node(
            "LinearClassifier",
            ["X"],
            ["I", "Y"],
            domain="ai.onnx.ml",
            classlabels_ints=[1],
            coefficients=[-0.58, -0.29, -0.09],
            intercepts=[2.7],
            multi_class=0,
            post_transform=post,
        )
        graph = make_graph([node1], "ml", [X], [In, Y])
        model = make_model_gen_version(graph, opset_imports=OPSETS)
        onnx.checker.check_model(model)
        # onnxruntime answer seems odd.
        # self._check_ort(model, {"X": x}, rev=True)
        sess = ReferenceEvaluator(model)
        got = sess.run(None, {"X": x})
        assert_allclose(got[1], expected[1], atol=1e-6)
        assert_allclose(got[0], expected[0])

    @staticmethod
    def _get_test_tree_ensemble_opset_latest(
        aggregate_function,
        rule=Mode.LEQ,
        unique_targets=False,
        input_type=TensorProto.FLOAT,
    ):
        X = make_tensor_value_info("X", input_type, [None, None])
        Y = make_tensor_value_info("Y", input_type, [None, None])
        if unique_targets:
            weights = [
                1.0,
                10.0,
                100.0,
                1000.0,
                10000.0,
                100000.0,
            ]
        else:
            weights = [
                0.07692307978868484,
                0.5,
                0.5,
                0.0,
                0.2857142984867096,
                0.5,
            ]
        node = make_node(
            "TreeEnsemble",
            ["X"],
            ["Y"],
            domain="ai.onnx.ml",
            n_targets=1,
            aggregate_function=aggregate_function,
            membership_values=None,
            nodes_missing_value_tracks_true=None,
            nodes_hitrates=None,
            post_transform=0,
            tree_roots=[0, 2],
            nodes_splits=make_tensor(
                "node_splits",
                input_type,
                (4,),
                [
                    0.26645058393478394,
                    0.6214364767074585,
                    -0.5592705607414246,
                    -0.7208403944969177,
                ],
            ),
            nodes_featureids=[0, 2, 0, 0],
            nodes_modes=make_tensor(
                "nodes_modes",
                TensorProto.UINT8,
                (4,),
                [rule] * 4,
            ),
            nodes_truenodeids=[1, 0, 3, 4],
            nodes_trueleafs=[0, 1, 1, 1],
            nodes_falsenodeids=[2, 1, 3, 5],
            nodes_falseleafs=[1, 1, 0, 1],
            leaf_targetids=[0, 0, 0, 0, 0, 0],
            leaf_weights=make_tensor(
                "leaf_weights", input_type, (len(weights),), weights
            ),
        )
        graph = make_graph([node], "ml", [X], [Y])
        return make_model_gen_version(graph, opset_imports=OPSETS)

    @staticmethod
    def _get_test_tree_ensemble_regressor(
        aggregate_function, rule="BRANCH_LEQ", unique_targets=False, base_values=None
    ):
        opsets = [make_opsetid("", TARGET_OPSET), make_opsetid("ai.onnx.ml", 3)]
        X = make_tensor_value_info("X", TensorProto.FLOAT, [None, None])
        Y = make_tensor_value_info("Y", TensorProto.FLOAT, [None, None])
        if unique_targets:
            targets = [
                1.0,
                10.0,
                100.0,
                1000.0,
                10000.0,
                100000.0,
            ]
        else:
            targets = [
                0.07692307978868484,
                0.5,
                0.5,
                0.0,
                0.2857142984867096,
                0.5,
            ]
        node1 = make_node(
            "TreeEnsembleRegressor",
            ["X"],
            ["Y"],
            domain="ai.onnx.ml",
            n_targets=1,
            aggregate_function=aggregate_function,
            base_values=base_values,
            nodes_falsenodeids=[4, 3, 0, 0, 0, 2, 0, 4, 0, 0],
            nodes_featureids=[0, 2, 0, 0, 0, 0, 0, 2, 0, 0],
            nodes_hitrates=[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
            nodes_missing_value_tracks_true=[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            nodes_modes=[
                rule,
                rule,
                "LEAF",
                "LEAF",
                "LEAF",
                rule,
                "LEAF",
                rule,
                "LEAF",
                "LEAF",
            ],
            nodes_nodeids=[0, 1, 2, 3, 4, 0, 1, 2, 3, 4],
            nodes_treeids=[0, 0, 0, 0, 0, 1, 1, 1, 1, 1],
            nodes_truenodeids=[1, 2, 0, 0, 0, 1, 0, 3, 0, 0],
            nodes_values=[
                0.26645058393478394,
                0.6214364767074585,
                0.0,
                0.0,
                0.0,
                -0.7208403944969177,
                0.0,
                -0.5592705607414246,
                0.0,
                0.0,
            ],
            post_transform="NONE",
            target_ids=[0, 0, 0, 0, 0, 0],
            target_nodeids=[2, 3, 4, 1, 3, 4],
            target_treeids=[0, 0, 0, 1, 1, 1],
            target_weights=targets,
        )
        graph = make_graph([node1], "ml", [X], [Y])
        model = make_model_gen_version(graph, opset_imports=opsets)
        onnx.checker.check_model(model)
        return model

    @pytest.mark.parametrize(
        "aggregate_function, expected_result, opset5",
        tuple(
            itertools.chain.from_iterable(
                (
                    (
                        AggregationFunction.SUM if opset5 else "SUM",
                        np.array(
                            [[0.576923], [0.576923], [0.576923]], dtype=np.float32
                        ),
                        opset5,
                    ),
                    (
                        AggregationFunction.AVERAGE if opset5 else "AVERAGE",
                        np.array(
                            [[0.288462], [0.288462], [0.288462]], dtype=np.float32
                        ),
                        opset5,
                    ),
                    (
                        AggregationFunction.MIN if opset5 else "MIN",
                        np.array(
                            [[0.076923], [0.076923], [0.076923]], dtype=np.float32
                        ),
                        opset5,
                    ),
                    (
                        AggregationFunction.MAX if opset5 else "MAX",
                        np.array([[0.5], [0.5], [0.5]], dtype=np.float32),
                        opset5,
                    ),
                )
                for opset5 in [True, False]
            )
        ),
    )
    @pytest.mark.skipif(not ONNX_ML, reason="onnx not compiled with ai.onnx.ml")
    def test_tree_ensemble_regressor_aggregation_functions(
        self, aggregate_function, expected_result, opset5
    ):
        x = np.arange(9).reshape((-1, 3)).astype(np.float32) / 10 - 0.5
        model_factory = (
            self._get_test_tree_ensemble_opset_latest
            if opset5
            else self._get_test_tree_ensemble_regressor
        )
        model_proto = model_factory(
            aggregate_function,
        )
        sess = ReferenceEvaluator(model_proto)
        (actual,) = sess.run(None, {"X": x})
        assert_allclose(actual, expected_result, atol=1e-6)

    @pytest.mark.parametrize(
        "rule, expected, opset5",
        tuple(
            itertools.chain.from_iterable(
                (
                    (
                        Mode.LEQ if opset5 else "BRANCH_LEQ",
                        np.array(
                            [[0.576923], [0.576923], [0.576923]], dtype=np.float32
                        ),
                        opset5,
                    ),
                    (
                        Mode.GT if opset5 else "BRANCH_GT",
                        np.array([[0.5], [0.5], [0.5]], dtype=np.float32),
                        opset5,
                    ),
                    (
                        Mode.LT if opset5 else "BRANCH_LT",
                        np.array(
                            [[0.576923], [0.576923], [0.576923]], dtype=np.float32
                        ),
                        opset5,
                    ),
                    (
                        Mode.GTE if opset5 else "BRANCH_GTE",
                        np.array([[0.5], [0.5], [0.5]], dtype=np.float32),
                        opset5,
                    ),
                    (
                        Mode.EQ if opset5 else "BRANCH_EQ",
                        np.array([[1.0], [1.0], [1.0]], dtype=np.float32),
                        opset5,
                    ),
                    (
                        Mode.NEQ if opset5 else "BRANCH_NEQ",
                        np.array(
                            [[0.076923], [0.076923], [0.076923]], dtype=np.float32
                        ),
                        opset5,
                    ),
                )
                for opset5 in [True, False]
            )
        ),
    )
    @pytest.mark.skipif(not ONNX_ML, reason="onnx not compiled with ai.onnx.ml")
    def test_tree_ensemble_regressor_rule(self, rule, expected, opset5):
        x = np.arange(9).reshape((-1, 3)).astype(np.float32) / 10 - 0.5
        model_factory = (
            self._get_test_tree_ensemble_opset_latest
            if opset5
            else self._get_test_tree_ensemble_regressor
        )
        aggregate_function = AggregationFunction.SUM if opset5 else "SUM"

        model_proto = model_factory(aggregate_function, rule)
        sess = ReferenceEvaluator(model_proto)
        (actual,) = sess.run(None, {"X": x})
        assert_allclose(actual, expected, atol=1e-6)

    @pytest.mark.skipif(not ONNX_ML, reason="onnx not compiled with ai.onnx.ml")
    def test_tree_ensemble_regressor_2_targets_opset3(self):
        X = make_tensor_value_info("X", TensorProto.FLOAT, [None, None])
        Y = make_tensor_value_info("Y", TensorProto.FLOAT, [None, None])
        opsets = [make_opsetid("", TARGET_OPSET), make_opsetid("ai.onnx.ml", 3)]
        node1 = make_node(
            "TreeEnsembleRegressor",
            ["X"],
            ["Y"],
            domain="ai.onnx.ml",
            n_targets=2,
            nodes_falsenodeids=[4, 3, 0, 0, 6, 0, 0, 4, 3, 0, 0, 6, 0, 0],
            nodes_featureids=[0, 2, 0, 0, 2, 0, 0, 0, 1, 0, 0, 2, 0, 0],
            nodes_hitrates=[
                1.0,
                1.0,
                1.0,
                1.0,
                1.0,
                1.0,
                1.0,
                1.0,
                1.0,
                1.0,
                1.0,
                1.0,
                1.0,
                1.0,
            ],
            nodes_missing_value_tracks_true=[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            nodes_modes=[
                "BRANCH_LEQ",
                "BRANCH_LEQ",
                "LEAF",
                "LEAF",
                "BRANCH_LEQ",
                "LEAF",
                "LEAF",
                "BRANCH_LEQ",
                "BRANCH_LEQ",
                "LEAF",
                "LEAF",
                "BRANCH_LEQ",
                "LEAF",
                "LEAF",
            ],
            nodes_nodeids=[0, 1, 2, 3, 4, 5, 6, 0, 1, 2, 3, 4, 5, 6],
            nodes_treeids=[0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1],
            nodes_truenodeids=[1, 2, 0, 0, 5, 0, 0, 1, 2, 0, 0, 5, 0, 0],
            nodes_values=[
                -0.3367232382297516,
                1.5326381921768188,
                0.0,
                0.0,
                -0.24646544456481934,
                0.0,
                0.0,
                -0.3367232382297516,
                0.6671845316886902,
                0.0,
                0.0,
                -0.24646544456481934,
                0.0,
                0.0,
            ],
            post_transform="NONE",
            target_ids=[0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1],
            target_nodeids=[2, 2, 3, 3, 5, 5, 6, 6, 2, 2, 3, 3, 5, 5, 6, 6],
            target_treeids=[0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1],
            target_weights=[
                0.0,
                2.5,
                0.5,
                3.0,
                0.15000000596046448,
                2.6500000953674316,
                0.5,
                3.0,
                0.02777777798473835,
                2.527777671813965,
                0.5,
                3.0,
                0.20000000298023224,
                2.700000047683716,
                0.5,
                3.0,
            ],
        )
        graph = make_graph([node1], "ml", [X], [Y])
        model = make_model_gen_version(graph, opset_imports=opsets)
        onnx.checker.check_model(model)
        x = np.arange(9).reshape((-1, 3)).astype(np.float32) / 10 - 0.5
        expected = np.array(
            [[0.027778, 5.027778], [1.0, 6.0], [1.0, 6.0]], dtype=np.float32
        )
        self._check_ort(model, {"X": x}, equal=True)
        sess = ReferenceEvaluator(model)
        got = sess.run(None, {"X": x})
        assert_allclose(got[0], expected, atol=1e-6)

    @pytest.mark.skipif(not ONNX_ML, reason="onnx not compiled with ai.onnx.ml")
    def test_tree_ensemble_regressor_missing_opset3(self):
        x = np.arange(9).reshape((-1, 3)).astype(np.float32) / 10 - 0.5
        x[2, 0] = 5
        x[1, :] = np.nan
        expected = np.array([[100001.0], [100100.0], [100100.0]], dtype=np.float32)
        model = self._get_test_tree_ensemble_regressor("SUM", unique_targets=True)
        self._check_ort(model, {"X": x}, equal=True)
        sess = ReferenceEvaluator(model)
        got = sess.run(None, {"X": x})
        assert_allclose(got[0], expected, atol=1e-6)
        assert "op_type=TreeEnsembleRegressor" in str(sess.rt_nodes_[0])

    @pytest.mark.skipif(not ONNX_ML, reason="onnx not compiled with ai.onnx.ml")
    @pytest.mark.parametrize("input_type", [TensorProto.FLOAT, TensorProto.DOUBLE])
    def test_tree_ensemble_missing_opset5(self, input_type):
        model = self._get_test_tree_ensemble_opset_latest(
            AggregationFunction.SUM, Mode.LEQ, True, input_type
        )
        np_dtype = onnx.helper.tensor_dtype_to_np_dtype(input_type)
        x = np.arange(9).reshape((-1, 3)).astype(np_dtype) / 10 - 0.5
        x[2, 0] = 5
        x[1, :] = np.nan
        expected = np.array([[100001.0], [100100.0], [100100.0]], dtype=np_dtype)
        session = ReferenceEvaluator(model)
        (actual,) = session.run(None, {"X": x})
        assert_allclose(actual, expected, atol=1e-6)

    @pytest.mark.skipif(not ONNX_ML, reason="onnx not compiled with ai.onnx.ml")
    def test_tree_ensemble_regressor_missing_opset5_float16(self):
        model = self._get_test_tree_ensemble_opset_latest(
            AggregationFunction.SUM, Mode.LEQ, False, TensorProto.FLOAT16
        )
        np_dtype = np.float16
        x = np.arange(9).reshape((-1, 3)).astype(np_dtype) / 10 - 0.5
        x[2, 0] = 5
        x[1, :] = np.nan
        expected = np.array([[0.577], [1.0], [1.0]], dtype=np_dtype)
        session = ReferenceEvaluator(model)
        (actual,) = session.run(None, {"X": x})
        assert_allclose(actual, expected, atol=1e-6)

    @pytest.mark.skipif(not ONNX_ML, reason="onnx not compiled with ai.onnx.ml")
    def test_single_tree_ensemble(self):
        X = make_tensor_value_info("X", TensorProto.DOUBLE, [None, None])
        Y = make_tensor_value_info("Y", TensorProto.DOUBLE, [None, None])
        node = make_node(
            "TreeEnsemble",
            ["X"],
            ["Y"],
            domain="ai.onnx.ml",
            n_targets=2,
            membership_values=None,
            nodes_missing_value_tracks_true=None,
            nodes_hitrates=None,
            aggregate_function=1,
            post_transform=PostTransform.NONE,
            tree_roots=[0],
            nodes_modes=make_tensor(
                "nodes_modes",
                TensorProto.UINT8,
                (3,),
                [Mode.LEQ] * 3,
            ),
            nodes_featureids=[0, 0, 0],
            nodes_splits=make_tensor(
                "nodes_splits",
                TensorProto.DOUBLE,
                (3,),
                np.array([3.14, 1.2, 4.2], dtype=np.float64),
            ),
            nodes_truenodeids=[1, 0, 1],
            nodes_trueleafs=[0, 1, 1],
            nodes_falsenodeids=[2, 2, 3],
            nodes_falseleafs=[0, 1, 1],
            leaf_targetids=[0, 1, 0, 1],
            leaf_weights=make_tensor(
                "leaf_weights",
                TensorProto.DOUBLE,
                (4,),
                np.array([5.23, 12.12, -12.23, 7.21], dtype=np.float64),
            ),
        )
        graph = make_graph([node], "ml", [X], [Y])
        model = make_model_gen_version(
            graph,
            opset_imports=[
                make_opsetid("", TARGET_OPSET),
                make_opsetid("ai.onnx.ml", 5),
            ],
        )
        onnx.checker.check_model(model)
        session = ReferenceEvaluator(model)
        (output,) = session.run(
            None,
            {
                "X": np.array([1.2, 3.4, -0.12, 1.66, 4.14, 1.77], np.float64).reshape(
                    3, 2
                )
            },
        )
        np.testing.assert_equal(
            output, np.array([[5.23, 0], [5.23, 0], [0, 12.12]], dtype=np.float64)
        )

    @pytest.mark.skipif(not ONNX_ML, reason="onnx not compiled with ai.onnx.ml")
    def test_tree_ensemble_regressor_set_membership_opset5(self):
        X = make_tensor_value_info("X", TensorProto.FLOAT, [None, None])
        Y = make_tensor_value_info("Y", TensorProto.FLOAT, [None, None])
        node = make_node(
            "TreeEnsemble",
            ["X"],
            ["Y"],
            domain="ai.onnx.ml",
            n_targets=4,
            aggregate_function=AggregationFunction.SUM,
            membership_values=make_tensor(
                "membership_values",
                TensorProto.FLOAT,
                (8,),
                [1.2, 3.7, 8, 9, np.nan, 12, 7, np.nan],
            ),
            nodes_missing_value_tracks_true=None,
            nodes_hitrates=None,
            post_transform=PostTransform.NONE,
            tree_roots=[0],
            nodes_modes=make_tensor(
                "nodes_modes",
                TensorProto.UINT8,
                (3,),
                [Mode.LEQ, Mode.MEMBER, Mode.MEMBER],
            ),
            nodes_featureids=[0, 0, 0],
            nodes_splits=make_tensor(
                "nodes_splits",
                TensorProto.FLOAT,
                (3,),
                np.array([11, 232344.0, np.nan], dtype=np.float32),
            ),
            nodes_trueleafs=[0, 1, 1],
            nodes_truenodeids=[1, 0, 1],
            nodes_falseleafs=[1, 0, 1],
            nodes_falsenodeids=[2, 2, 3],
            leaf_targetids=[0, 1, 2, 3],
            leaf_weights=make_tensor(
                "leaf_weights", TensorProto.FLOAT, (4,), [1, 10, 1000, 100]
            ),
        )
        graph = make_graph([node], "ml", [X], [Y])
        model = make_model_gen_version(
            graph,
            opset_imports=OPSETS,
        )
        onnx.checker.check_model(model)
        session = ReferenceEvaluator(model)
        X = np.array([1.2, 3.4, -0.12, np.nan, 12, 7], np.float32).reshape(-1, 1)
        expected = np.array(
            [
                [1, 0, 0, 0],
                [0, 0, 0, 100],
                [0, 0, 0, 100],
                [0, 0, 1000, 0],
                [0, 0, 1000, 0],
                [0, 10, 0, 0],
            ],
            dtype=np.float32,
        )
        (output,) = session.run(None, {"X": X})
        np.testing.assert_equal(output, expected)

    @staticmethod
    def _get_test_svm_regressor(kernel_type, kernel_params):
        X = make_tensor_value_info("X", TensorProto.FLOAT, [None, None])
        Y = make_tensor_value_info("Y", TensorProto.FLOAT, [None, None])
        node1 = make_node(
            "SVMRegressor",
            ["X"],
            ["Y"],
            domain="ai.onnx.ml",
            coefficients=[
                1.0,
                -1.0,
                0.8386201858520508,
                -0.8386201858520508,
                0.4470679759979248,
                -1.0,
                0.5529320240020752,
            ],
            kernel_params=kernel_params,
            kernel_type=kernel_type,
            n_supports=7,
            post_transform="NONE",
            rho=[0.5460880398750305],
            support_vectors=[
                -0.12850627303123474,
                0.08915442228317261,
                0.06881910562515259,
                -0.07938569784164429,
                -0.22557435929775238,
                -0.26520243287086487,
                0.9246066212654114,
                -0.025557516142725945,
                -0.5900523662567139,
                0.9735698699951172,
                -1.3385062217712402,
                0.3393094539642334,
                0.9432410001754761,
                -0.5228781700134277,
                0.5557093620300293,
                0.4191802740097046,
                0.43368014693260193,
                -1.0569839477539062,
                2.3318440914154053,
                0.06202844902873039,
                -0.9502395987510681,
            ],
        )
        graph = make_graph([node1], "ml", [X], [Y])
        model = make_model_gen_version(graph, opset_imports=OPSETS)
        onnx.checker.check_model(model)
        return model

    @pytest.mark.parametrize(
        "kernel, params, expected",
        [
            (
                "LINEAR",
                [0.42438405752182007, 0.0, 3.0],
                np.array([[-0.468206], [0.227487], [0.92318]], dtype=np.float32),
            ),
            (
                "POLY",
                [0.3426632285118103, 0.0, 3.0],
                np.array([[0.527084], [0.543578], [0.546506]], dtype=np.float32),
            ),
            (
                "RBF",
                [0.30286383628845215, 0.0, 3.0],
                np.array([[0.295655], [0.477876], [0.695292]], dtype=np.float32),
            ),
            (
                "SIGMOID",
                [0.30682486295700073, 0.0, 3.0],
                np.array([[0.239304], [0.448929], [0.661689]], dtype=np.float32),
            ),
        ],
    )
    @pytest.mark.skipif(not ONNX_ML, reason="onnx not compiled with ai.onnx.ml")
    def test_svm_regressor(self, kernel, params, expected):
        x = np.arange(9).reshape((-1, 3)).astype(np.float32) / 10 - 0.5

        model = self._get_test_svm_regressor(kernel, params)
        self._check_ort(model, {"X": x}, atol=1e-6)
        sess = ReferenceEvaluator(model)
        (got,) = sess.run(None, {"X": x})
        assert_allclose(got, expected, atol=1e-6)

    @staticmethod
    def _get_test_tree_ensemble_classifier_binary(post_transform):
        X = make_tensor_value_info("X", TensorProto.FLOAT, [None, None])
        In = make_tensor_value_info("I", TensorProto.INT64, [None])
        Y = make_tensor_value_info("Y", TensorProto.FLOAT, [None, None])
        node1 = make_node(
            "TreeEnsembleClassifier",
            ["X"],
            ["I", "Y"],
            domain="ai.onnx.ml",
            class_ids=[0, 0, 0, 0, 0, 0, 0],
            class_nodeids=[2, 3, 5, 6, 1, 3, 4],
            class_treeids=[0, 0, 0, 0, 1, 1, 1],
            class_weights=[
                0.0,
                0.1764705926179886,
                0.0,
                0.5,
                0.0,
                0.0,
                0.4285714328289032,
            ],
            classlabels_int64s=[0, 1],
            nodes_falsenodeids=[4, 3, 0, 0, 6, 0, 0, 2, 0, 4, 0, 0],
            nodes_featureids=[2, 2, 0, 0, 1, 0, 0, 2, 0, 0, 0, 0],
            nodes_hitrates=[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
            nodes_missing_value_tracks_true=[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            nodes_modes=[
                "BRANCH_LEQ",
                "BRANCH_LEQ",
                "LEAF",
                "LEAF",
                "BRANCH_LEQ",
                "LEAF",
                "LEAF",
                "BRANCH_LEQ",
                "LEAF",
                "BRANCH_LEQ",
                "LEAF",
                "LEAF",
            ],
            nodes_nodeids=[0, 1, 2, 3, 4, 5, 6, 0, 1, 2, 3, 4],
            nodes_treeids=[0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1],
            nodes_truenodeids=[1, 2, 0, 0, 5, 0, 0, 1, 0, 3, 0, 0],
            nodes_values=[
                0.6874135732650757,
                -0.3654803931713104,
                0.0,
                0.0,
                -1.926770806312561,
                0.0,
                0.0,
                -0.3654803931713104,
                0.0,
                -2.0783839225769043,
                0.0,
                0.0,
            ],
            post_transform=post_transform,
        )
        graph = make_graph([node1], "ml", [X], [In, Y])
        model = make_model_gen_version(
            graph,
            opset_imports=[
                make_opsetid("", TARGET_OPSET),
                make_opsetid("ai.onnx.ml", 3),
            ],
        )
        onnx.checker.check_model(model)
        return model

    @pytest.mark.parametrize(
        "post, expected",
        [
            (
                "NONE",
                (
                    np.array([0, 1, 1], dtype=np.int64),
                    np.array(
                        [[1.0, 0.0], [0.394958, 0.605042], [0.394958, 0.605042]],
                        dtype=np.float32,
                    ),
                ),
            ),
            (
                "LOGISTIC",
                (
                    np.array([0, 1, 1], dtype=np.int64),
                    np.array(
                        [[0.5, 0.5], [0.353191, 0.646809], [0.353191, 0.646809]],
                        dtype=np.float32,
                    ),
                ),
            ),
            (
                "SOFTMAX",
                (
                    np.array([0, 1, 1], dtype=np.int64),
                    np.array(
                        [[0.5, 0.5], [0.229686, 0.770314], [0.229686, 0.770314]],
                        dtype=np.float32,
                    ),
                ),
            ),
            (
                "SOFTMAX_ZERO",
                (
                    np.array([0, 1, 1], dtype=np.int64),
                    np.array(
                        [[0.5, 0.5], [0.229686, 0.770314], [0.229686, 0.770314]],
                        dtype=np.float32,
                    ),
                ),
            ),
            (
                "PROBIT",
                (
                    np.array([0, 1, 1], dtype=np.int64),
                    np.array(
                        [[0.0, 0.0], [-0.266426, 0.266426], [-0.266426, 0.266426]],
                        dtype=np.float32,
                    ),
                ),
            ),
        ],
    )
    @pytest.mark.skipif(not ONNX_ML, reason="onnx not compiled with ai.onnx.ml")
    def test_tree_ensemble_classifier_binary(self, post, expected):
        x = (np.arange(9).reshape((-1, 3)) - 5).astype(np.float32) / 5
        model = self._get_test_tree_ensemble_classifier_binary(post)
        if post == "NONE":
            self._check_ort(model, {"X": x})
        sess = ReferenceEvaluator(model)
        got = sess.run(None, {"X": x})
        assert_allclose(got[1], expected[1], atol=1e-6)
        assert_allclose(got[0], expected[0])

    @staticmethod
    def _get_test_tree_ensemble_classifier_multi(post_transform):
        X = make_tensor_value_info("X", TensorProto.FLOAT, [None, None])
        In = make_tensor_value_info("I", TensorProto.INT64, [None])
        Y = make_tensor_value_info("Y", TensorProto.FLOAT, [None, None])
        node1 = make_node(
            "TreeEnsembleClassifier",
            ["X"],
            ["I", "Y"],
            domain="ai.onnx.ml",
            class_ids=[0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2],
            class_nodeids=[2, 2, 2, 3, 3, 3, 4, 4, 4, 1, 1, 1, 3, 3, 3, 4, 4, 4],
            class_treeids=[0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1],
            class_weights=[
                0.46666666865348816,
                0.0,
                0.03333333507180214,
                0.20000000298023224,
                0.23999999463558197,
                0.05999999865889549,
                0.0,
                0.5,
                0.0,
                0.5,
                0.0,
                0.0,
                0.44999998807907104,
                0.0,
                0.05000000074505806,
                0.10294117778539658,
                0.19117647409439087,
                0.20588235557079315,
            ],
            classlabels_int64s=[0, 1, 2],
            nodes_falsenodeids=[4, 3, 0, 0, 0, 2, 0, 4, 0, 0],
            nodes_featureids=[1, 0, 0, 0, 0, 1, 0, 0, 0, 0],
            nodes_hitrates=[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
            nodes_missing_value_tracks_true=[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            nodes_modes=[
                "BRANCH_LEQ",
                "BRANCH_LEQ",
                "LEAF",
                "LEAF",
                "LEAF",
                "BRANCH_LEQ",
                "LEAF",
                "BRANCH_LEQ",
                "LEAF",
                "LEAF",
            ],
            nodes_nodeids=[0, 1, 2, 3, 4, 0, 1, 2, 3, 4],
            nodes_treeids=[0, 0, 0, 0, 0, 1, 1, 1, 1, 1],
            nodes_truenodeids=[1, 2, 0, 0, 0, 1, 0, 3, 0, 0],
            nodes_values=[
                1.2495747804641724,
                -0.3050493597984314,
                0.0,
                0.0,
                0.0,
                -1.6830512285232544,
                0.0,
                -0.6751254796981812,
                0.0,
                0.0,
            ],
            post_transform=post_transform,
        )
        graph = make_graph([node1], "ml", [X], [In, Y])
        model = make_model_gen_version(
            graph,
            opset_imports=[
                make_opsetid("", TARGET_OPSET),
                make_opsetid("ai.onnx.ml", 3),
            ],
        )
        onnx.checker.check_model(model)
        return model

    @pytest.mark.parametrize(
        "post, expected",
        [
            (
                "NONE",
                (
                    np.array([0, 0, 1], dtype=np.int64),
                    np.array(
                        [
                            [0.916667, 0.0, 0.083333],
                            [0.569608, 0.191176, 0.239216],
                            [0.302941, 0.431176, 0.265882],
                        ],
                        dtype=np.float32,
                    ),
                ),
            ),
            (
                "LOGISTIC",
                (
                    np.array([0, 0, 1], dtype=np.int64),
                    np.array(
                        [
                            [0.714362, 0.5, 0.520821],
                            [0.638673, 0.547649, 0.55952],
                            [0.575161, 0.606155, 0.566082],
                        ],
                        dtype=np.float32,
                    ),
                ),
            ),
            (
                "SOFTMAX",
                (
                    np.array([0, 0, 1], dtype=np.int64),
                    np.array(
                        [
                            [0.545123, 0.217967, 0.23691],
                            [0.416047, 0.284965, 0.298988],
                            [0.322535, 0.366664, 0.310801],
                        ],
                        dtype=np.float32,
                    ),
                ),
            ),
            (
                "SOFTMAX_ZERO",
                (
                    np.array([0, 0, 1], dtype=np.int64),
                    np.array(
                        [
                            [0.697059, 0.0, 0.302941],
                            [0.416047, 0.284965, 0.298988],
                            [0.322535, 0.366664, 0.310801],
                        ],
                        dtype=np.float32,
                    ),
                ),
            ),
            (
                "PROBIT",
                (
                    np.array([0, 0, 1], dtype=np.int64),
                    np.array(
                        [
                            [1.383104, 0, -1.383105],
                            [0.175378, -0.873713, -0.708922],
                            [-0.516003, -0.173382, -0.625385],
                        ],
                        dtype=np.float32,
                    ),
                ),
            ),
        ],
    )
    @pytest.mark.skipif(not ONNX_ML, reason="onnx not compiled with ai.onnx.ml")
    def test_tree_ensemble_classifier_multi(self, post, expected):
        x = (np.arange(9).reshape((-1, 3)) - 5).astype(np.float32) / 5
        model = self._get_test_tree_ensemble_classifier_multi(post)
        if post != "PROBIT":
            self._check_ort(model, {"X": x}, atol=1e-5)
        sess = ReferenceEvaluator(model)
        got = sess.run(None, {"X": x})
        assert_allclose(got[1], expected[1], atol=1e-6)
        assert_allclose(got[0], expected[0])

    @staticmethod
    def _get_test_svm_classifier_binary(post_transform, probability=True, linear=False):
        X = make_tensor_value_info("X", TensorProto.FLOAT, [None, None])
        In = make_tensor_value_info("I", TensorProto.INT64, [None])
        Y = make_tensor_value_info("Y", TensorProto.FLOAT, [None, None])
        if linear:
            kwargs = {
                "classlabels_ints": [0, 1, 2, 3],
                "coefficients": [
                    -1.55181212e-01,
                    2.42698956e-01,
                    7.01893432e-03,
                    4.07614474e-01,
                    -3.24927823e-02,
                    2.79897536e-04,
                    -1.95771302e-01,
                    -3.52437368e-01,
                    -2.15973096e-02,
                    -4.38190277e-01,
                    4.56869105e-02,
                    -1.29375499e-02,
                ],
                "kernel_params": [0.001, 0.0, 3.0],
                "kernel_type": "LINEAR",
                "prob_a": [-5.139118194580078],
                "prob_b": [0.06399919837713242],
                "rho": [-0.07489691, -0.1764396, -0.21167431, -0.51619097],
                "post_transform": post_transform,
            }
        else:
            kwargs = {
                "classlabels_ints": [0, 1],
                "coefficients": [1.0, 1.0, 1.0, 1.0, -1.0, -1.0, -1.0, -1.0],
                "kernel_params": [0.3824487328529358, 0.0, 3.0],
                "kernel_type": "RBF",
                "prob_a": [-5.139118194580078],
                "prob_b": [0.06399919837713242],
                "rho": [0.16708599030971527],
                "support_vectors": [
                    0.19125767052173615,
                    -1.062204122543335,
                    0.5006636381149292,
                    -0.5892484784126282,
                    -0.3196830451488495,
                    0.0984845906496048,
                    0.24746321141719818,
                    -1.1535362005233765,
                    0.4109955430030823,
                    -0.5937694907188416,
                    -1.3183348178863525,
                    -1.6423596143722534,
                    0.558641254901886,
                    -0.9218668341636658,
                    0.6264089345932007,
                    -0.16060839593410492,
                    -0.6365169882774353,
                    0.8335472345352173,
                    0.7539799213409424,
                    -0.3970031440258026,
                    -0.1780400276184082,
                    -0.616622805595398,
                    0.49261474609375,
                    0.4470972716808319,
                ],
                "vectors_per_class": [4, 4],
                "post_transform": post_transform,
            }

        if not probability:
            del kwargs["prob_a"]
            del kwargs["prob_b"]
        node1 = make_node(
            "SVMClassifier", ["X"], ["I", "Y"], domain="ai.onnx.ml", **kwargs
        )
        graph = make_graph([node1], "ml", [X], [In, Y])
        model = make_model_gen_version(graph, opset_imports=OPSETS)
        onnx.checker.check_model(model)
        return model

    @pytest.mark.parametrize(
        "post, expected",
        [
            (
                "NONE",
                (
                    np.array([0, 1, 1], dtype=np.int64),
                    np.array(
                        [
                            [0.993287, 0.006713],
                            [0.469401, 0.530599],
                            [0.014997, 0.985003],
                        ],
                        dtype=np.float32,
                    ),
                ),
            ),
            (
                "LOGISTIC",
                (
                    np.array([0, 1, 1], dtype=np.int64),
                    np.array(
                        [
                            [0.729737, 0.501678],
                            [0.615242, 0.629623],
                            [0.503749, 0.7281],
                        ],
                        dtype=np.float32,
                    ),
                ),
            ),
            (
                "SOFTMAX",
                (
                    np.array([0, 1, 1], dtype=np.int64),
                    np.array(
                        [
                            [0.728411, 0.271589],
                            [0.484705, 0.515295],
                            [0.274879, 0.725121],
                        ],
                        dtype=np.float32,
                    ),
                ),
            ),
            (
                "SOFTMAX_ZERO",
                (
                    np.array([0, 1, 1], dtype=np.int64),
                    np.array(
                        [
                            [0.728411, 0.271589],
                            [0.484705, 0.515295],
                            [0.274879, 0.725121],
                        ],
                        dtype=np.float32,
                    ),
                ),
            ),
            (
                "PROBIT",
                (
                    np.array([0, 1, 1], dtype=np.int64),
                    np.array(
                        [
                            [2.469393, -2.469391],
                            [-0.076776, 0.076776],
                            [-2.16853, 2.16853],
                        ],
                        dtype=np.float32,
                    ),
                ),
            ),
        ],
    )
    @pytest.mark.skipif(not ONNX_ML, reason="onnx not compiled with ai.onnx.ml")
    def test_svm_classifier_binary(self, post, expected):
        x = (np.arange(9).reshape((-1, 3)) - 5).astype(np.float32) / 5
        model = self._get_test_svm_classifier_binary(post)
        self._check_ort(model, {"X": x}, rev=True, atol=1e-5)
        sess = ReferenceEvaluator(model)
        got = sess.run(None, {"X": x})
        assert_allclose(got[1], expected[1], atol=1e-5)
        assert_allclose(got[0], expected[0])

    @pytest.mark.parametrize(
        "post, expected",
        [
            (
                "NONE",
                (
                    np.array([0, 1, 1], dtype=np.int64),
                    np.array(
                        [
                            [-0.986073, 0.986073],
                            [0.011387, -0.011387],
                            [0.801808, -0.801808],
                        ],
                        dtype=np.float32,
                    ),
                ),
            ),
            (
                "LOGISTIC",
                (
                    np.array([0, 1, 1], dtype=np.int64),
                    np.array(
                        [
                            [0.271688, 0.728312],
                            [0.502847, 0.497153],
                            [0.690361, 0.309639],
                        ],
                        dtype=np.float32,
                    ),
                ),
            ),
            (
                "SOFTMAX",
                (
                    np.array([0, 1, 1], dtype=np.int64),
                    np.array(
                        [
                            [0.122158, 0.877842],
                            [0.505693, 0.494307],
                            [0.832523, 0.167477],
                        ],
                        dtype=np.float32,
                    ),
                ),
            ),
            (
                "SOFTMAX_ZERO",
                (
                    np.array([0, 1, 1], dtype=np.int64),
                    np.array(
                        [
                            [0.122158, 0.877842],
                            [0.505693, 0.494307],
                            [0.832523, 0.167477],
                        ],
                        dtype=np.float32,
                    ),
                ),
            ),
        ],
    )
    @pytest.mark.skipif(not ONNX_ML, reason="onnx not compiled with ai.onnx.ml")
    def test_svm_classifier_binary_noprob(self, post, expected):
        x = (np.arange(9).reshape((-1, 3)) - 5).astype(np.float32) / 5
        model = self._get_test_svm_classifier_binary(post, probability=False)
        if post not in {"LOGISTIC", "SOFTMAX", "SOFTMAX_ZERO", "PROBIT"}:
            self._check_ort(model, {"X": x}, rev=True, atol=1e-5)
        sess = ReferenceEvaluator(model)
        got = sess.run(None, {"X": x})
        assert_allclose(got[1], expected[1], atol=1e-6)
        assert_allclose(got[0], expected[0])

    @pytest.mark.parametrize(
        "post, expected",
        [
            (
                "NONE",
                (
                    np.array([2, 3, 0], dtype=np.int64),
                    np.array(
                        [
                            [-0.118086, -0.456685, 0.415783, 0.334506],
                            [-0.061364, -0.231444, 0.073899, 0.091242],
                            [-0.004642, -0.006203, -0.267985, -0.152023],
                        ],
                        dtype=np.float32,
                    ),
                ),
            ),
            (
                "LOGISTIC",
                (
                    np.array([2, 3, 0], dtype=np.int64),
                    np.array(
                        [
                            [0.470513, 0.387773, 0.602474, 0.582855],
                            [0.484664, 0.442396, 0.518466, 0.522795],
                            [0.498839, 0.498449, 0.433402, 0.462067],
                        ],
                        dtype=np.float32,
                    ),
                ),
            ),
            (
                "SOFTMAX",
                (
                    np.array([2, 3, 0], dtype=np.int64),
                    np.array(
                        [
                            [0.200374, 0.14282, 0.341741, 0.315065],
                            [0.240772, 0.203115, 0.275645, 0.280467],
                            [0.275491, 0.275061, 0.211709, 0.237739],
                        ],
                        dtype=np.float32,
                    ),
                ),
            ),
            (
                "SOFTMAX_ZERO",
                (
                    np.array([2, 3, 0], dtype=np.int64),
                    np.array(
                        [
                            [0.200374, 0.14282, 0.341741, 0.315065],
                            [0.240772, 0.203115, 0.275645, 0.280467],
                            [0.275491, 0.275061, 0.211709, 0.237739],
                        ],
                        dtype=np.float32,
                    ),
                ),
            ),
            (
                "PROBIT",
                (
                    np.array([2, 3, 0], dtype=np.int64),
                    np.array(
                        [
                            [np.nan, np.nan, -0.212698, -0.427529],
                            [np.nan, np.nan, -1.447414, -1.333286],
                            [np.nan, np.nan, np.nan, np.nan],
                        ],
                        dtype=np.float32,
                    ),
                ),
            ),
        ],
    )
    @pytest.mark.skipif(not ONNX_ML, reason="onnx not compiled with ai.onnx.ml")
    def test_svm_classifier_noprob_linear(self, post, expected):
        x = (np.arange(9).reshape((-1, 3)) - 5).astype(np.float32) / 5
        model = self._get_test_svm_classifier_binary(
            post, probability=False, linear=True
        )
        self._check_ort(model, {"X": x}, rev=True, atol=1e-5)
        sess = ReferenceEvaluator(model)
        got = sess.run(None, {"X": x})
        assert_allclose(got[1], expected[1], atol=1e-6)
        assert_allclose(got[0], expected[0])

    @pytest.mark.parametrize(
        "post, expected",
        [
            (
                "NONE",
                (
                    np.array([2, 3, 0], dtype=np.int64),
                    np.array(
                        [
                            [-0.118086, -0.456685, 0.415783, 0.334506],
                            [-0.061364, -0.231444, 0.073899, 0.091242],
                            [-0.004642, -0.006203, -0.267985, -0.152023],
                        ],
                        dtype=np.float32,
                    ),
                ),
            ),
            (
                "LOGISTIC",
                (
                    np.array([2, 3, 0], dtype=np.int64),
                    np.array(
                        [
                            [0.470513, 0.387773, 0.602474, 0.582855],
                            [0.484664, 0.442396, 0.518466, 0.522795],
                            [0.498839, 0.498449, 0.433402, 0.462067],
                        ],
                        dtype=np.float32,
                    ),
                ),
            ),
            (
                "SOFTMAX",
                (
                    np.array([2, 3, 0], dtype=np.int64),
                    np.array(
                        [
                            [0.200374, 0.14282, 0.341741, 0.315065],
                            [0.240772, 0.203115, 0.275645, 0.280467],
                            [0.275491, 0.275061, 0.211709, 0.237739],
                        ],
                        dtype=np.float32,
                    ),
                ),
            ),
            (
                "SOFTMAX_ZERO",
                (
                    np.array([2, 3, 0], dtype=np.int64),
                    np.array(
                        [
                            [0.200374, 0.14282, 0.341741, 0.315065],
                            [0.240772, 0.203115, 0.275645, 0.280467],
                            [0.275491, 0.275061, 0.211709, 0.237739],
                        ],
                        dtype=np.float32,
                    ),
                ),
            ),
            (
                "PROBIT",
                (
                    np.array([2, 3, 0], dtype=np.int64),
                    np.array(
                        [
                            [np.nan, np.nan, -0.212698, -0.427529],
                            [np.nan, np.nan, -1.447414, -1.333286],
                            [np.nan, np.nan, np.nan, np.nan],
                        ],
                        dtype=np.float32,
                    ),
                ),
            ),
        ],
    )
    @pytest.mark.skipif(not ONNX_ML, reason="onnx not compiled with ai.onnx.ml")
    def test_svm_classifier_linear(self, post, expected):
        # prob_a, prob_b are not used in this case.
        x = (np.arange(9).reshape((-1, 3)) - 5).astype(np.float32) / 5
        model = self._get_test_svm_classifier_binary(
            post, probability=True, linear=True
        )
        self._check_ort(model, {"X": x}, rev=True, atol=1e-5)
        sess = ReferenceEvaluator(model)
        got = sess.run(None, {"X": x})
        assert_allclose(got[1], expected[1], atol=1e-6)
        assert_allclose(got[0], expected[0])

    @staticmethod
    def _get_test_svm_classifier_linear_sv(post_transform, probability=True):
        X = make_tensor_value_info("X", TensorProto.FLOAT, [None, None])
        In = make_tensor_value_info("I", TensorProto.INT64, [None])
        Y = make_tensor_value_info("Y", TensorProto.FLOAT, [None, None])
        kwargs = {
            "classlabels_ints": [0, 1],
            "coefficients": [
                0.766398549079895,
                0.0871576070785522,
                0.110420741140842,
                -0.963976919651031,
            ],
            "support_vectors": [
                4.80000019073486,
                3.40000009536743,
                1.89999997615814,
                5.0,
                3.0,
                1.60000002384186,
                4.5,
                2.29999995231628,
                1.29999995231628,
                5.09999990463257,
                2.5,
                3.0,
            ],
            "kernel_params": [0.122462183237076, 0.0, 3.0],
            "kernel_type": "LINEAR",
            "prob_a": [-5.139118194580078],
            "prob_b": [0.06399919837713242],
            "rho": [2.23510527610779],
            "post_transform": post_transform,
            "vectors_per_class": [3, 1],
        }

        if not probability:
            del kwargs["prob_a"]
            del kwargs["prob_b"]
        node1 = make_node(
            "SVMClassifier", ["X"], ["I", "Y"], domain="ai.onnx.ml", **kwargs
        )
        graph = make_graph([node1], "ml", [X], [In, Y])
        model = make_model_gen_version(graph, opset_imports=OPSETS)
        onnx.checker.check_model(model)
        return model

    @pytest.mark.parametrize(
        "post, expected",
        [
            (
                "NONE",
                (
                    np.array([0, 0, 0], dtype=np.int64),
                    np.array(
                        [
                            [-2.662655, 2.662655],
                            [-2.21481, 2.21481],
                            [-1.766964, 1.766964],
                        ],
                        dtype=np.float32,
                    ),
                ),
            ),
            (
                "LOGISTIC",
                (
                    np.array([0, 0, 0], dtype=np.int64),
                    np.array(
                        [
                            [0.065213, 0.934787],
                            [0.098428, 0.901572],
                            [0.14592, 0.85408],
                        ],
                        dtype=np.float32,
                    ),
                ),
            ),
            (
                "SOFTMAX",
                (
                    np.array([0, 0, 0], dtype=np.int64),
                    np.array(
                        [
                            [0.004843, 0.995157],
                            [0.011779, 0.988221],
                            [0.028362, 0.971638],
                        ],
                        dtype=np.float32,
                    ),
                ),
            ),
            (
                "SOFTMAX_ZERO",
                (
                    np.array([0, 0, 0], dtype=np.int64),
                    np.array(
                        [
                            [0.004843, 0.995157],
                            [0.011779, 0.988221],
                            [0.028362, 0.971638],
                        ],
                        dtype=np.float32,
                    ),
                ),
            ),
        ],
    )
    @pytest.mark.skipif(not ONNX_ML, reason="onnx not compiled with ai.onnx.ml")
    def test_svm_classifier_binary_noprob_linear_sv(self, post, expected):
        x = (np.arange(9).reshape((-1, 3)) - 5).astype(np.float32) / 5

        model = self._get_test_svm_classifier_linear_sv(post, probability=False)
        if post not in {"LOGISTIC", "SOFTMAX", "SOFTMAX_ZERO"}:
            self._check_ort(model, {"X": x}, rev=True, atol=1e-5)
        sess = ReferenceEvaluator(model)
        got = sess.run(None, {"X": x})
        assert_allclose(got[1], expected[1], atol=1e-6)
        assert_allclose(got[0], expected[0])

    @staticmethod
    def _get_test_svm_regressor_linear(post_transform, one_class=0):
        X = make_tensor_value_info("X", TensorProto.FLOAT, [None, None])
        Y = make_tensor_value_info("Y", TensorProto.FLOAT, [None])
        kwargs = {
            "coefficients": [0.28290501, -0.0266512, 0.01674867],
            "kernel_params": [0.001, 0.0, 3.0],
            "kernel_type": "LINEAR",
            "rho": [1.24032312],
            "post_transform": post_transform,
            "n_supports": 0,
            "one_class": one_class,
        }

        node1 = make_node("SVMRegressor", ["X"], ["Y"], domain="ai.onnx.ml", **kwargs)
        graph = make_graph([node1], "ml", [X], [Y])
        model = make_model_gen_version(graph, opset_imports=OPSETS)
        onnx.checker.check_model(model)
        return model

    @pytest.mark.parametrize(
        "post, expected",
        [
            (
                "NONE",
                np.array(
                    [[0.96869], [1.132491], [1.296293]],
                    dtype=np.float32,
                ),
            ),
        ],
    )
    @pytest.mark.skipif(not ONNX_ML, reason="onnx not compiled with ai.onnx.ml")
    def test_svm_regressor_linear(self, post, expected):
        x = (np.arange(9).reshape((-1, 3)) - 5).astype(np.float32) / 5
        model = self._get_test_svm_regressor_linear(post)
        self._check_ort(model, {"X": x}, rev=True, atol=1e-5)
        sess = ReferenceEvaluator(model)
        (got,) = sess.run(None, {"X": x})
        assert_allclose(got, expected, atol=1e-6)

    @pytest.mark.parametrize(
        "post, expected",
        [
            (
                "NONE",
                np.array(
                    [[1.0], [1.0], [1.0]],
                    dtype=np.float32,
                ),
            )
        ],
    )
    @pytest.mark.skipif(not ONNX_ML, reason="onnx not compiled with ai.onnx.ml")
    def test_svm_regressor_linear_one_class(self, post, expected):
        x = (np.arange(9).reshape((-1, 3)) - 5).astype(np.float32) / 5
        model = self._get_test_svm_regressor_linear(post, one_class=1)
        self._check_ort(model, {"X": x}, rev=True, atol=1e-5)
        sess = ReferenceEvaluator(model)
        (got,) = sess.run(None, {"X": x})
        assert_allclose(got, expected, atol=1e-6)

    def test_onnxrt_tfidf_vectorizer_ints(self):
        inputi = np.array([[1, 1, 3, 3, 3, 7], [8, 6, 7, 5, 6, 8]]).astype(np.int64)
        output = np.array(
            [[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 1.0]]
        ).astype(np.float32)

        ngram_counts = np.array([0, 4]).astype(np.int64)
        ngram_indexes = np.array([0, 1, 2, 3, 4, 5, 6]).astype(np.int64)
        pool_int64s = np.array([2, 3, 5, 4, 5, 6, 7, 8, 6, 7]).astype(  # unigrams
            np.int64
        )  # bigrams

        model = make_model_gen_version(
            make_graph(
                [
                    make_node(
                        "TfIdfVectorizer",
                        ["tokens"],
                        ["out"],
                        mode="TF",
                        min_gram_length=2,
                        max_gram_length=2,
                        max_skip_count=0,
                        ngram_counts=ngram_counts,
                        ngram_indexes=ngram_indexes,
                        pool_int64s=pool_int64s,
                    )
                ],
                "tfidf",
                [make_tensor_value_info("tokens", TensorProto.INT64, [None, None])],
                [make_tensor_value_info("out", TensorProto.FLOAT, [None, None])],
            ),
            opset_imports=OPSETS,
        )

        oinf = ReferenceEvaluator(model)
        res = oinf.run(None, {"tokens": inputi})
        assert output.tolist() == res[0].tolist()

    def test_onnxrt_tfidf_vectorizer_strings(self):
        inputi = np.array(
            [["i1", "i1", "i3", "i3", "i3", "i7"], ["i8", "i6", "i7", "i5", "i6", "i8"]]
        )
        output = np.array(
            [[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 1.0]]
        ).astype(np.float32)

        ngram_counts = np.array([0, 4]).astype(np.int64)
        ngram_indexes = np.array([0, 1, 2, 3, 4, 5, 6]).astype(np.int64)
        pool_strings = np.array(
            ["i2", "i3", "i5", "i4", "i5", "i6", "i7", "i8", "i6", "i7"]
        )

        model = make_model_gen_version(
            make_graph(
                [
                    make_node(
                        "TfIdfVectorizer",
                        ["tokens"],
                        ["out"],
                        mode="TF",
                        min_gram_length=2,
                        max_gram_length=2,
                        max_skip_count=0,
                        ngram_counts=ngram_counts,
                        ngram_indexes=ngram_indexes,
                        pool_strings=pool_strings,
                    )
                ],
                "tfidf",
                [make_tensor_value_info("tokens", TensorProto.STRING, [None, None])],
                [make_tensor_value_info("out", TensorProto.FLOAT, [None, None])],
            ),
            opset_imports=OPSETS,
        )

        oinf = ReferenceEvaluator(model)
        res = oinf.run(None, {"tokens": inputi})
        assert output.tolist() == res[0].tolist()