Repository: onnx/sklearn-onnx
Pull request #1165 (2025-02-04 13:30:53)

Fix KNNImputer

ashtarimo: KNNImputer converter doesn't generate the same values for the imputed missing values. I am using:

onnxruntim 1.18.1, onnx  1.17.0, and skl2onnx 1.17.0

Here is the code:


import numpy as np
import onnx
import skl2onnx
from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType
from sklearn.impute import KNNImputer

np.random.seed(42)
data = np.random.randn(100, 10)
data[99, 9] = np.nan
imputer = KNNImputer(n_neighbors= 5)
imputer.fit(data)
dataft = imputer.transform(data)
initial_type = [('float_input', FloatTensorType([None, data.shape[1]]))]
onnx_model = convert_sklearn(imputer, initial_types=initial_type)
with open("knn_imputer.onnx", "wb") as f:
    f.write(onnx_model.SerializeToString())
sess = rt.InferenceSession("knn_imputer.onnx")
input_data = data.astype(np.float32)
input_name = sess.get_inputs()[0].name
output_name = sess.get_outputs()[0].name
res = sess.run([output_name], {input_name: input_data})

The KNNIMputer generates value of -0.67440104 for the missing value, and the onnx model, generates -1.0071397 for the same missing value

Patch:
diff --git a/skl2onnx/common/_container.py b/skl2onnx/common/_container.py
index 9ca6a1a..45caea4 100644
--- a/skl2onnx/common/_container.py
+++ b/skl2onnx/common/_container.py
@@ -11,11 +11,18 @@ import numpy as np
 from scipy.sparse import coo_matrix
 from onnx import SparseTensorProto, ValueInfoProto
 from onnx.defs import onnx_opset_version, get_all_schemas_with_history
-from onnx.helper import make_node, make_tensor, make_attribute, make_sparse_tensor
+from onnx.helper import (
+    make_node,
+    make_tensor,
+    make_attribute,
+    make_sparse_tensor,
+    make_tensor_value_info,
+)
 import onnx.onnx_cpp2py_export.defs as C
 from onnxconverter_common.onnx_ops import __dict__ as dict_apply_operation
 from ..proto import TensorProto
 from .utils import get_domain
+from .graph_builder_opset import Opset
 
 
 logger = getLogger("skl2onnx")
@@ -222,6 +229,7 @@ class ModelComponentContainer(_WhiteBlackContainer):
         white_op=None,
         black_op=None,
         verbose=0,
+        as_function=False,
     ):
         """
         :param target_opset: number, for example, 7 for *ONNX 1.2*, and
@@ -233,6 +241,7 @@ class ModelComponentContainer(_WhiteBlackContainer):
         :param black_op: black list of ONNX nodes allowed
             while converting a pipeline, if empty, none are blacklisted
         :param verbose: display information while converting
+        :param as_function: to export as a local function
         """
         _WhiteBlackContainer.__init__(
             self, white_op=white_op, black_op=black_op, verbose=verbose
@@ -251,6 +260,7 @@ class ModelComponentContainer(_WhiteBlackContainer):
         # ONNX nodes (type: NodeProto) used to define computation
         # structure
         self.nodes = []
+        self.as_function = as_function
         # ONNX operators' domain-version pair set. They will be added
         # into opset_import field in the final ONNX model.
         self.node_domain_version_pair_sets = set()
@@ -266,6 +276,14 @@ class ModelComponentContainer(_WhiteBlackContainer):
         self.options = options
         # All registered models.
         self.registered_models = registered_models
+        self.local_functions = {}
+
+    @property
+    def main_opset(self):
+        return self.target_opset
+
+    def make_local_function(self, domain: str, name: str, container, optimize=False):
+        self.local_functions[domain, name] = container
 
     def swap_names(self, old_name, new_name):
         """
@@ -389,6 +407,14 @@ class ModelComponentContainer(_WhiteBlackContainer):
         """
         self.inputs.append(self._make_value_info(variable))
 
+    def make_tensor_input(self, name):
+        self.inputs.append(make_tensor_value_info(name, 0, None))
+        return name
+
+    def make_tensor_output(self, name):
+        self.outputs.append(make_tensor_value_info(name, 0, None))
+        return name
+
     def add_output(self, variable):
         """
         Adds our *Variable* object defined *_parser.py* into the the
@@ -621,6 +647,7 @@ class ModelComponentContainer(_WhiteBlackContainer):
                       attributes' names and attributes' values,
                       respectively.
         """
+        assert op_type != "knn_imputer_column" or op_domain
         if "axes" in attrs and (
             attrs["axes"] is None or not isinstance(attrs["axes"], (list, np.ndarray))
         ):
@@ -650,6 +677,12 @@ class ModelComponentContainer(_WhiteBlackContainer):
             ",".join(outputs),
             name,
         )
+        if not hasattr(self, "_added_names_"):
+            self._added_names_ = set()
+        assert all(
+            n not in self._added_names_ for n in outputs
+        ), f"One output node in {outputs} was already added, added={self._added_names_}"
+        self._added_names_ |= set(outputs)
         try:
             common = set(inputs) & set(outputs)
         except TypeError as e:
@@ -763,7 +796,7 @@ class ModelComponentContainer(_WhiteBlackContainer):
         key = domain, op_type
         vers = self._op_versions.get(key, None)
         if vers is None:
-            if domain == "com.microsoft":
+            if domain in ("com.microsoft", "local_domain"):
                 # avoid a not necessarily necessary warning
                 vers = 1
             else:
@@ -875,6 +908,7 @@ class ModelComponentContainer(_WhiteBlackContainer):
         n_iter = 0
         missing_ops = []
         cont = True
+        parent = {}
         while cont and n_iter < len(self.nodes) * 2:
             n_iter += 1
             missing_names = set()
@@ -900,11 +934,16 @@ class ModelComponentContainer(_WhiteBlackContainer):
                 order[key] = maxi
                 maxi += 1
                 for name in node.output:
+                    if name not in parent:
+                        parent[name] = []
+                    parent[name].append(node)
                     if name in order:
                         raise RuntimeError(
-                            "Unable to sort a node (cycle). An output was "
-                            "already ordered with name %r (iteration=%r)."
-                            "" % (name, n_iter)
+                            f"Unable to sort a node (cycle). An output was "
+                            f"already ordered with name {name!r} (iteration={n_iter})\n"
+                            f"--\n{[f'{id(n)}-{n.op_type}' for n in parent[name]]}\n"
+                            f"--\n{[n.op_type for n in self.nodes]}\n"
+                            f"--\n{pprint.pformat(order)}"
                         )
                     order[name] = maxi
             if len(missing_names) == 0:
@@ -1030,3 +1069,6 @@ class ModelComponentContainer(_WhiteBlackContainer):
         )
         map_nodes = {str(id(node)): node for node in self.nodes}
         self.nodes = [map_nodes[_[-1]] for _ in topo]
+
+    def get_op_builder(self, scope):
+        return Opset(self, scope)
diff --git a/skl2onnx/common/_topology.py b/skl2onnx/common/_topology.py
index 5d44d3c..7b17d23 100644
--- a/skl2onnx/common/_topology.py
+++ b/skl2onnx/common/_topology.py
@@ -8,7 +8,13 @@ from logging import getLogger
 from collections import OrderedDict
 import numpy as np
 from onnx import onnx_pb as onnx_proto
-from onnx.helper import make_graph, make_model, make_tensor_value_info
+from onnx.helper import (
+    make_graph,
+    make_model,
+    make_tensor_value_info,
+    make_node,
+    make_function,
+)
 from onnxconverter_common.data_types import (
     DataType,
     TensorType,
@@ -1463,69 +1469,16 @@ class Topology:
             container.add_output(o)
 
 
-def convert_topology(
-    topology,
-    model_name,
-    doc_string,
-    target_opset,
-    options=None,
+def make_model_from_container(
+    container: ModelComponentContainer,
     remove_identity=True,
-    verbose=0,
+    doc_string="",
+    verbose: int = 0,
+    onnx_target_opset=None,
+    model_name="",
+    as_function=False,
+    domain_name="",
 ):
-    """
-    This function is used to convert our Topology object defined in
-    _parser.py into a ONNX model (type: ModelProto).
-
-    :param topology: The Topology object we are going to convert
-    :param model_name: GraphProto's name. Let "model" denote the
-        returned model. The string "model_name" would be
-        assigned to "model.graph.name."
-    :param doc_string: A string attached to the produced model
-    :param target_opset: number or dictionary,
-        for example, 7 for ONNX 1.2, and 8 for ONNX 1.3,
-        a dictionary is used to indicate different opset for
-        different domains
-    :param options: see :ref:`l-conv-options`
-    :param remove_identity: removes identity nodes
-        include '1.1.2', '1.2', and so on.
-    :param verbose: displays information while converting
-    :return: a ONNX ModelProto
-    """
-    if target_opset is None:
-        target_opset = get_latest_tested_opset_version()
-    if isinstance(target_opset, dict):
-        onnx_target_opset = target_opset.get("", get_latest_tested_opset_version())
-    else:
-        onnx_target_opset = target_opset
-    if onnx_target_opset > get_opset_number_from_onnx():
-        found = get_opset_number_from_onnx()
-        raise RuntimeError(
-            "Parameter target_opset {} > {} is higher than the "
-            "version of the installed onnx package. See "
-            "https://github.com/onnx/onnx/blob/main/docs/"
-            "Versioning.md#released-versions"
-            ".".format(onnx_target_opset, found)
-        )
-    if onnx_target_opset > get_latest_tested_opset_version():
-        warnings.warn(
-            "Parameter target_opset {} > {} is higher than the "
-            "the latest tested version"
-            ".".format(onnx_target_opset, get_latest_tested_opset_version()),
-            stacklevel=0,
-        )
-
-    container = ModelComponentContainer(
-        target_opset,
-        options=options,
-        registered_models=topology.registered_models,
-        white_op=topology.raw_model._white_op,
-        black_op=topology.raw_model._black_op,
-        verbose=verbose,
-    )
-
-    # Traverse the graph from roots to leaves
-    # This loop could eventually be parallelized.
-    topology.convert_operators(container=container, verbose=verbose)
     container.ensure_topological_order()
 
     if len(container.inputs) == 0:
@@ -1542,6 +1495,7 @@ def convert_topology(
 
     # Create a graph from its main components
     if container.target_opset_onnx < 9:
+        assert not as_function, "Functions not supported for opset < 9."
         # When calling ModelComponentContainer's add_initializer(...),
         # nothing is added into the input list. However, for ONNX target
         # opset < 9, initializers should also be a part of model's
@@ -1578,19 +1532,37 @@ def convert_topology(
         # In ONNX opset 9 and above, initializers are included as
         # operator inputs and therefore do not need to be passed as
         # extra_inputs.
-        graph = make_graph(
-            container.nodes,
-            model_name,
-            container.inputs,
-            container.outputs,
-            container.initializers,
-        )
-
-    # Add extra information related to the graph
-    graph.value_info.extend(container.value_info)
+        if as_function:
+            graph = make_function(
+                domain_name,
+                model_name,
+                [i.name for i in container.inputs],
+                [i.name for i in container.outputs],
+                [
+                    *[
+                        make_node("Constant", [], [init.name], value=init)
+                        for init in container.initializers
+                    ],
+                    *container.nodes,
+                ],
+                [],  # opsets
+            )
+        else:
+            graph = make_graph(
+                container.nodes,
+                model_name,
+                container.inputs,
+                container.outputs,
+                container.initializers,
+            )
 
     # Create model
-    onnx_model = make_model(graph)
+    if as_function:
+        onnx_model = graph
+    else:
+        # Add extra information related to the graph
+        graph.value_info.extend(container.value_info)
+        onnx_model = make_model(graph)
 
     # Update domain version
     opv = min(
@@ -1605,12 +1577,13 @@ def convert_topology(
             print("[convert_topology] +opset: name=%r, version=%s" % ("", opv))
 
     # Add extra information
-    irv = OPSET_TO_IR_VERSION.get(opv, onnx_proto.IR_VERSION)
-    onnx_model.ir_version = irv
-    onnx_model.producer_name = utils.get_producer()
-    onnx_model.producer_version = utils.get_producer_version()
-    onnx_model.domain = utils.get_domain()
-    onnx_model.model_version = utils.get_model_version()
+    if not as_function:
+        irv = OPSET_TO_IR_VERSION.get(opv, onnx_proto.IR_VERSION)
+        onnx_model.ir_version = irv
+        onnx_model.producer_name = utils.get_producer()
+        onnx_model.producer_version = utils.get_producer_version()
+        onnx_model.domain = utils.get_domain()
+        onnx_model.model_version = utils.get_model_version()
     onnx_model.doc_string = doc_string
 
     # Removes many identity nodes,
@@ -1621,9 +1594,95 @@ def convert_topology(
     if remove_identity:
         onnx_model = onnx_remove_node_identity(onnx_model)
 
+    # process functions
+    for key, cont in container.local_functions.items():
+        domain, name = key
+        proto = make_model_from_container(
+            cont,
+            remove_identity=remove_identity,
+            verbose=verbose,
+            as_function=True,
+            model_name=name,
+            domain_name=domain,
+            onnx_target_opset=container.target_opset_onnx,
+        )
+        onnx_model.functions.append(proto)
+
     return onnx_model
 
 
+def convert_topology(
+    topology,
+    model_name,
+    doc_string,
+    target_opset,
+    options=None,
+    remove_identity=True,
+    verbose=0,
+):
+    """
+    This function is used to convert our Topology object defined in
+    _parser.py into a ONNX model (type: ModelProto).
+
+    :param topology: The Topology object we are going to convert
+    :param model_name: GraphProto's name. Let "model" denote the
+        returned model. The string "model_name" would be
+        assigned to "model.graph.name."
+    :param doc_string: A string attached to the produced model
+    :param target_opset: number or dictionary,
+        for example, 7 for ONNX 1.2, and 8 for ONNX 1.3,
+        a dictionary is used to indicate different opset for
+        different domains
+    :param options: see :ref:`l-conv-options`
+    :param remove_identity: removes identity nodes
+        include '1.1.2', '1.2', and so on.
+    :param verbose: displays information while converting
+    :return: a ONNX ModelProto
+    """
+    if target_opset is None:
+        target_opset = get_latest_tested_opset_version()
+    if isinstance(target_opset, dict):
+        onnx_target_opset = target_opset.get("", get_latest_tested_opset_version())
+    else:
+        onnx_target_opset = target_opset
+    if onnx_target_opset > get_opset_number_from_onnx():
+        found = get_opset_number_from_onnx()
+        raise RuntimeError(
+            "Parameter target_opset {} > {} is higher than the "
+            "version of the installed onnx package. See "
+            "https://github.com/onnx/onnx/blob/main/docs/"
+            "Versioning.md#released-versions"
+            ".".format(onnx_target_opset, found)
+        )
+    if onnx_target_opset > get_latest_tested_opset_version():
+        warnings.warn(
+            "Parameter target_opset {} > {} is higher than the "
+            "the latest tested version"
+            ".".format(onnx_target_opset, get_latest_tested_opset_version()),
+            stacklevel=0,
+        )
+
+    container = ModelComponentContainer(
+        target_opset,
+        options=options,
+        registered_models=topology.registered_models,
+        white_op=topology.raw_model._white_op,
+        black_op=topology.raw_model._black_op,
+        verbose=verbose,
+    )
+    # Traverse the graph from roots to leaves
+    # This loop could eventually be parallelized.
+    topology.convert_operators(container=container, verbose=verbose)
+    return make_model_from_container(
+        container,
+        remove_identity=remove_identity,
+        doc_string=doc_string,
+        verbose=verbose,
+        onnx_target_opset=onnx_target_opset,
+        model_name=model_name,
+    )
+
+
 def _update_domain_version(container, onnx_model, verbose=0):
     # Merge operator sets for the same domain, the largest version
     # number would be kept
@@ -1637,10 +1696,12 @@ def _update_domain_version(container, onnx_model, verbose=0):
             )
 
     # Fill operator sets
+    done = set()
     i = 0
     for op_domain, op_version in purified_operator_set.items():
         if op_version is None:
             continue
+        done.add(op_domain)
         if i == 0 and len(onnx_model.opset_import) == 1:
             # Overwrite the default operator set created by
             # make_model(...)
@@ -1672,6 +1733,14 @@ def _update_domain_version(container, onnx_model, verbose=0):
                 "this model, which requires at least opset "
                 "%d." % (container.target_opset_any_domain(op_domain), op_version)
             )
+    for k, v in container.target_opset_all.items():
+        if k in done:
+            continue
+        # Some opsets may be forgotten if they appear in subgraphs.
+        op_set = onnx_model.opset_import.add()
+        op_set.domain = k
+        op_set.version = v
+
     return "" in purified_operator_set
 
 
diff --git a/skl2onnx/common/graph_builder_opset.py b/skl2onnx/common/graph_builder_opset.py
new file mode 100644
index 0000000..44c2462
--- /dev/null
+++ b/skl2onnx/common/graph_builder_opset.py
@@ -0,0 +1,275 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from functools import partial
+from typing import List, Optional, Union
+import numpy as np
+from onnx.helper import np_dtype_to_tensor_dtype
+
+
+class Opset:
+    """
+    Makes it easier to write onnx graph.
+    The method name is the node type.
+
+    :param graph_builder: the builder or container
+    :param allow_unknown: allows unknown operators, otherwise,
+        fails this class does not the expected number of outputs
+    """
+
+    # defined for opset >= 18
+    # name: number of expected outputs
+    _implemented = {
+        "Abs": 1,
+        "Add": 1,
+        "And": 1,
+        "ArgMax": 1,
+        "ArgMin": 1,
+        "Cast": 1,
+        "CastLike": 1,
+        "Celu": 1,
+        "Concat": 1,
+        "Constant": 1,
+        "ConstantOfShape": 1,
+        "Cos": 1,
+        "Cosh": 1,
+        "Div": 1,
+        "Dropout": 2,
+        "Elu": 1,
+        "Equal": 1,
+        "Exp": 1,
+        "Expand": 1,
+        "Flatten": 1,
+        "Gather": 1,
+        "GatherElements": 1,
+        "GatherND": 1,
+        "Gemm": 1,
+        "Greater": 1,
+        "GreaterOrEqual": 1,
+        "Identity": 1,
+        "MatMul": 1,
+        "MaxPool": 2,
+        "Mul": 1,
+        "Less": 1,
+        "LessOrEqual": 1,
+        "Log": 1,
+        "LogSoftmax": 1,
+        "Neg": 1,
+        "Not": 1,
+        "Or": 1,
+        "Pow": 1,
+        "Range": 1,
+        "Reciprocal": 1,
+        "ReduceMax": 1,
+        "ReduceMean": 1,
+        "ReduceMin": 1,
+        "ReduceSum": 1,
+        "Relu": 1,
+        "Reshape": 1,
+        "ScatterElements": 1,
+        "ScatterND": 1,
+        "Shape": 1,
+        "Sigmoid": 1,
+        "Sin": 1,
+        "Sinh": 1,
+        "Slice": 1,
+        "Softmax": 1,
+        "Sqrt": 1,
+        "Squeeze": 1,
+        "Sub": 1,
+        "Tile": 1,
+        "TopK": 2,
+        "Transpose": 1,
+        "Trilu": 1,
+        "Unsqueeze": 1,
+        "Where": 1,
+    }
+
+    def __init__(
+        self,
+        container,
+        scope,
+        allow_unknown: bool = False,
+    ):
+        self.container = container
+        self.scope = scope
+        self.allow_unknown = allow_unknown
+
+    def __getattr__(self, name):
+        if name in self._implemented:
+            return partial(self.make_node, name)
+        if name in self.__dict__:
+            return self.__dict__[name]
+        return partial(self._make_node, name)
+
+    def _make_node(self, op_type, *args, outputs=None, **kwargs):
+        if outputs is None:
+            if op_type in self._implemented:
+                outputs = self._implemented[op_type]
+            elif op_type == "Split" and kwargs.get("domain", "") == "":
+                assert "num_outputs" in kwargs, (
+                    "Number of outputs is not implemented yet for operator "
+                    f"{op_type!r} and kwargs={kwargs}"
+                )
+                outputs = kwargs["num_outputs"]
+            else:
+                # We assume there is only one outputs.
+                outputs = 1
+        return self.make_node(
+            op_type, *args, outputs=outputs, allow_empty_shape=True, **kwargs
+        )
+
+    def make_node(
+        self,
+        op_type: str,
+        *inputs: Optional[Union[str, List[str]]],
+        outputs: Optional[Union[int, List[str], str]] = None,
+        domain: str = "",
+        name: Optional[str] = None,
+        allow_empty_shape: bool = False,
+        **kwargs,
+    ):
+        assert (
+            op_type != "Split" or outputs != 1
+        ), f"Operator Split is useless with one output, inputs={inputs}, outputs={outputs}"
+        if outputs is None:
+            num_outputs = self._implemented[op_type]
+            outputs = [
+                self.scope.get_unique_variable_name(f"_un_{op_type.lower()}")
+                for _ in range(num_outputs)
+            ]
+        elif isinstance(outputs, int):
+            outputs = [
+                self.scope.get_unique_variable_name(f"_un_{op_type.lower()}")
+                for _ in range(outputs)
+            ]
+        else:
+            assert all(
+                not self.scope.has_variable_name(n) for n in outputs
+            ), f"Duplicated names {outputs!r}"
+
+        if inputs is None:
+            inputs = []
+        assert (
+            op_type != "Reshape"
+            or len(inputs) != 2
+            or not isinstance(inputs[1], np.ndarray)
+            or inputs[1].dtype == np.int64
+        ), f"Suspicious shape {inputs[1]!r} for a Reshape"
+        new_inputs = []
+        for i in inputs:
+            assert not isinstance(
+                i, (list, tuple)
+            ), f"Wrong inputs for operator {op_type!r}: {inputs!r}"
+            if isinstance(i, str):
+                new_inputs.append(i)
+            elif hasattr(i, "name") and not hasattr(i, "detach"):
+                # torch.fx.Node
+                assert i.name is not None, f"Unexpected name for type {type(i)}"
+                new_inputs.append(i.name)
+            elif i is None:
+                # Optional input
+                new_inputs.append("")
+            elif isinstance(i, np.ndarray):
+                assert allow_empty_shape or 0 not in i.shape, (
+                    f"Not implemented for type(i)={type(i)}, i={i}, "
+                    f"inputs={inputs!r}, op_type={op_type!r}, i.shape={i.shape}"
+                    f""
+                )
+                cst_name = self.scope.get_unique_variable_name("cst")
+                self.container.add_initializer(
+                    cst_name,
+                    np_dtype_to_tensor_dtype(i.dtype),
+                    i.shape,
+                    list(i.ravel()),
+                )
+                new_inputs.append(cst_name)
+            else:
+                raise AssertionError(
+                    f"Not implemented for type(i)={type(i)}, i={i}, "
+                    f"inputs={inputs!r}, op_type={op_type!r}"
+                )
+
+        if self.allow_unknown and not self.container.get_opset(domain):
+            self.container.add_domain(domain)
+        self.container.add_node(
+            op_type,
+            new_inputs,
+            outputs,
+            op_domain=domain,
+            name=name or f"{self.__class__.__name__}",
+            **kwargs,
+        )
+        if len(outputs) == 1:
+            return outputs[0]
+        return outputs
+
+    @staticmethod
+    def _iaxes(op_type, axes) -> int:
+        if isinstance(axes, np.ndarray):
+            iaxes = axes.tolist()
+        elif isinstance(axes, int):
+            iaxes = [axes]
+        else:
+            raise RuntimeError(
+                f"Unable to call {op_type} on a dynamic input axis={axes}"
+            )
+        return iaxes
+
+    def ReduceMaxAnyOpset(self, *args, name: str = "ReduceMaxAnyOpset", **kwargs):
+        if len(args) == 1:
+            return self.ReduceMax(*args, name=name, **kwargs)
+        assert len(args) == 2, f"ReduceMaxAnyOpset expects 2 arguments not {len(args)}"
+        if self.container.main_opset >= 18:
+            return self.ReduceMax(*args, name=name, **kwargs)
+        return self.ReduceMax(
+            args[0], axes=self._iaxes("ReduceMax", args[1]), name=name, **kwargs
+        )
+
+    def ReduceMinAnyOpset(self, *args, name: str = "ReduceMinAnyOpset", **kwargs):
+        if len(args) == 1:
+            return self.ReduceMin(*args, name=name, **kwargs)
+        assert len(args) == 2, f"ReduceMaxAnyOpset expects 2 arguments not {len(args)}"
+        if self.container.main_opset >= 18:
+            return self.ReduceMin(*args, name=name, **kwargs)
+        return self.ReduceMin(
+            args[0], axes=self._iaxes("ReduceMin", args[1]), name=name, **kwargs
+        )
+
+    def ReduceMeanAnyOpset(self, *args, name: str = "ReduceMeanAnyOpset", **kwargs):
+        if len(args) == 1:
+            return self.ReduceMean(*args, name=name, **kwargs)
+        assert len(args) == 2, f"ReduceMeanAnyOpset expects 2 arguments not {len(args)}"
+        if self.container.main_opset >= 18:
+            return self.ReduceMean(*args, name=name, **kwargs)
+        return self.ReduceMean(
+            args[0], axes=self._iaxes("ReduceMean", args[1]), name=name, **kwargs
+        )
+
+    def ReduceSumAnyOpset(self, *args, name: str = "ReduceSumAnyOpset", **kwargs):
+        if len(args) == 1:
+            return self.ReduceSum(*args, name=name, **kwargs)
+        assert len(args) == 2, f"ReduceSumAnyOpset expects 2 arguments not {len(args)}"
+        if self.container.main_opset >= 13:
+            return self.ReduceSum(*args, name=name, **kwargs)
+        return self.ReduceSum(
+            args[0], axes=self._iaxes("ReduceSum", args[1]), name=name, **kwargs
+        )
+
+    def SqueezeAnyOpset(self, *args, name: str = "SqueezeAnyOpset", **kwargs):
+        if len(args) == 1 and len(kwargs) == 0:
+            return self.Squeeze(*args, name=name)
+        if self.container.main_opset >= 13 or len(args) == 1:
+            return self.Squeeze(*args, name=name, **kwargs)
+        return self.Squeeze(
+            args[0], axes=self._iaxes("Squeeze", args[1]), name=name, **kwargs
+        )
+
+    def UnsqueezeAnyOpset(self, *args, name: str = "UnsqueezeAnyOpset", **kwargs):
+        if len(args) == 1 and len(kwargs) == 0:
+            return self.Unsqueeze(*args, name=name)
+        assert len(args) == 2, f"UnsqueezeAnyOpset expects 2 arguments not {len(args)}"
+        if self.container.main_opset >= 13:
+            return self.Unsqueeze(*args, name=name, **kwargs)
+        return self.Unsqueeze(
+            args[0], axes=self._iaxes("Unsqueeze", args[1]), name=name, **kwargs
+        )
diff --git a/skl2onnx/common/onnx_optimisation_identity.py b/skl2onnx/common/onnx_optimisation_identity.py
index 804a226..d6f9a14 100644
--- a/skl2onnx/common/onnx_optimisation_identity.py
+++ b/skl2onnx/common/onnx_optimisation_identity.py
@@ -7,7 +7,7 @@ opset < 9.
 """
 
 from logging import getLogger
-from onnx.helper import make_graph
+from onnx.helper import make_graph, make_function
 from ._onnx_optimisation_common import (
     _rename_node_input,
     _rename_node_output,
@@ -48,8 +48,8 @@ def onnx_remove_node_identity(onnx_model, recursive=True, debug_info=None):
 
     graph = onnx_model
 
-    inputs = {i.name for i in graph.input}
-    outputs = {o.name for o in graph.output}
+    inputs = {(i if isinstance(i, str) else i.name) for i in graph.input}
+    outputs = {(o if isinstance(o, str) else o.name) for o in graph.output}
 
     def retrieve_idnodes(graph, existing_nodes):
         idnodes = []
@@ -160,13 +160,22 @@ def onnx_remove_node_identity(onnx_model, recursive=True, debug_info=None):
 
     # Finally create the new graph.
     nodes = list(filter(lambda n: n is not None, nodes))
-    graph = make_graph(
-        nodes,
-        onnx_model.name,
-        onnx_model.input,
-        onnx_model.output,
-        onnx_model.initializer,
-    )
-
-    graph.value_info.extend(onnx_model.value_info)
+    if hasattr(onnx_model, "initializer"):
+        graph = make_graph(
+            nodes,
+            onnx_model.name,
+            onnx_model.input,
+            onnx_model.output,
+            onnx_model.initializer,
+        )
+        graph.value_info.extend(onnx_model.value_info)
+    else:
+        graph = make_function(
+            onnx_model.domain,
+            onnx_model.name,
+            onnx_model.input,
+            onnx_model.output,
+            nodes,
+            onnx_model.opset_import,
+        )
     return graph
diff --git a/skl2onnx/operator_converters/nearest_neighbours.py b/skl2onnx/operator_converters/nearest_neighbours.py
index fadce52..1d09e1f 100644
--- a/skl2onnx/operator_converters/nearest_neighbours.py
+++ b/skl2onnx/operator_converters/nearest_neighbours.py
@@ -2,7 +2,15 @@
 
 
 import numpy as np
-from onnx.helper import make_tensor
+from onnx import TensorProto
+from onnx.helper import (
+    make_tensor,
+    tensor_dtype_to_np_dtype,
+    make_tensor_value_info,
+    make_graph,
+    make_node,
+)
+from onnx.numpy_helper import from_array
 from ..algebra.onnx_ops import (
     OnnxAdd,
     OnnxArgMax,
@@ -18,7 +26,6 @@ from ..algebra.onnx_ops import (
     OnnxMax,
     OnnxMul,
     OnnxNeg,
-    OnnxNot,
     OnnxReciprocal,
     OnnxReduceMeanApi18,
     OnnxReduceSumApi11,
@@ -193,7 +200,7 @@ def onnx_nearest_neighbors_indices_radius(
     zero = OnnxCast(
         OnnxConstantOfShape(shape, op_version=opv), op_version=opv, to=proto_dtype
     )
-    tensor_value = py_make_float_array(-1, dtype=np.float32, as_tensor=True)
+    tensor_value = py_make_float_array(-1, dtype=dtype, as_tensor=True)
     minus = OnnxCast(
         OnnxConstantOfShape(shape, op_version=opv, value=tensor_value),
         op_version=opv,
@@ -788,12 +795,839 @@ def _nearest_neighbours(
     return node[1], missing_input_name
 
 
+def make_dict_idx_map(g: ModelComponentContainer, scope: Scope, itype: int):
+    gr = ModelComponentContainer({"": g.main_opset}, as_function=True)
+    x = gr.make_tensor_input("x")
+    row_missing_idx = gr.make_tensor_input("row_missing_idx")
+    op = gr.get_op_builder(scope)
+    init7_s_0 = op.Constant(
+        value=from_array(np.array(0, dtype=np.int64), name="value"),
+        outputs=["init7_s_0"],
+    )
+    init7_s_1 = op.Constant(
+        value=from_array(np.array(1, dtype=np.int64), name="value"),
+        outputs=["init7_s_1"],
+    )
+    init7_s1__1 = op.Constant(
+        value=from_array(np.array([-1], dtype=np.int64), name="value"),
+        outputs=["init7_s1__1"],
+    )
+    _shape_x0 = op.Shape(x, end=1, start=0, outputs=["_shape_x0"])
+    _shape_row_missing_idx0 = op.Shape(
+        row_missing_idx, end=1, start=0, outputs=["_shape_row_missing_idx0"]
+    )
+    sym_size_int_4 = op.SqueezeAnyOpset(
+        _shape_row_missing_idx0, outputs=["sym_size_int_4"]
+    )
+    zeros = op.ConstantOfShape(
+        _shape_x0,
+        value=from_array(np.array([0], dtype=np.int64), name="value"),
+        outputs=["zeros"],
+    )
+    arange = op.Range(init7_s_0, sym_size_int_4, init7_s_1, outputs=["arange"])
+    _onx_unsqueeze_row_missing_idx0 = op.UnsqueezeAnyOpset(
+        row_missing_idx, init7_s1__1, outputs=["_onx_unsqueeze_row_missing_idx0"]
+    )
+    output_0 = op.ScatterND(
+        zeros, _onx_unsqueeze_row_missing_idx0, arange, outputs=["output_0"]
+    )
+    gr.make_tensor_output(output_0)
+    g.make_local_function(
+        container=gr, optimize=False, name="dict_idx_map", domain="local_domain"
+    )
+    return gr
+
+
+def make_dist_nan_euclidean(g: ModelComponentContainer, scope: Scope, itype: int):
+    gr = ModelComponentContainer({"": g.main_opset}, as_function=True)
+    x = gr.make_tensor_input("x")
+    y = gr.make_tensor_input("y")
+    op = gr.get_op_builder(scope)
+    c_lifted_tensor_0 = op.Constant(
+        value=from_array(
+            np.array(0.0, dtype=tensor_dtype_to_np_dtype(itype)), name="value"
+        ),
+        outputs=["c_lifted_tensor_0"],
+    )
+    c_lifted_tensor_1 = op.Constant(
+        value=from_array(
+            np.array(0.0, dtype=tensor_dtype_to_np_dtype(itype)), name="value"
+        ),
+        outputs=["c_lifted_tensor_1"],
+    )
+    c_lifted_tensor_2 = op.Constant(
+        value=from_array(
+            np.array(np.nan, dtype=tensor_dtype_to_np_dtype(itype)), name="value"
+        ),
+        outputs=["c_lifted_tensor_2"],
+    )
+    c_lifted_tensor_3 = op.Constant(
+        value=from_array(
+            np.array([1.0], dtype=tensor_dtype_to_np_dtype(itype)), name="value"
+        ),
+        outputs=["c_lifted_tensor_3"],
+    )
+    init1_s_ = op.Constant(
+        value=from_array(
+            np.array(-2.0, dtype=tensor_dtype_to_np_dtype(itype)), name="value"
+        ),
+        outputs=["init1_s_"],
+    )
+    init7_s1_1 = op.Constant(
+        value=from_array(np.array([1], dtype=np.int64), name="value"),
+        outputs=["init7_s1_1"],
+    )
+    init1_s1_ = op.Constant(
+        value=from_array(
+            np.array([0.0], dtype=tensor_dtype_to_np_dtype(itype)), name="value"
+        ),
+        outputs=["init1_s1_"],
+    )
+    init1_s_2 = op.Constant(
+        value=from_array(
+            np.array(1.0, dtype=tensor_dtype_to_np_dtype(itype)), name="value"
+        ),
+        outputs=["init1_s_2"],
+    )
+    init1_s_3 = op.Constant(
+        value=from_array(
+            np.array(0.0, dtype=tensor_dtype_to_np_dtype(itype)), name="value"
+        ),
+        outputs=["init1_s_3"],
+    )
+    init7_s2_1__1 = op.Constant(
+        value=from_array(np.array([1, -1], dtype=np.int64), name="value"),
+        outputs=["init7_s2_1__1"],
+    )
+    isnan = op.IsNaN(x, outputs=["isnan"])
+    _to_copy_2 = op.Cast(isnan, to=itype, outputs=["_to_copy_2"])
+    isnan_1 = op.IsNaN(y, outputs=["isnan_1"])
+    index_put = op.Where(isnan, c_lifted_tensor_0, x, outputs=["index_put"])
+    index_put_1 = op.Where(isnan_1, c_lifted_tensor_1, y, outputs=["index_put_1"])
+    mul_18 = op.Mul(index_put, index_put, outputs=["mul_18"])
+    mul_21 = op.Mul(index_put_1, index_put_1, outputs=["mul_21"])
+    matmul_2 = op.Gemm(_to_copy_2, mul_21, transA=0, transB=1, outputs=["matmul_2"])
+    _reshape_init1_s_0 = op.Reshape(
+        init1_s_, init7_s1_1, outputs=["_reshape_init1_s_0"]
+    )
+    index_put__2 = op.Mul(index_put, _reshape_init1_s_0, outputs=["index_put__2"])
+    matmul = op.Gemm(index_put__2, index_put_1, transA=0, transB=1, outputs=["matmul"])
+    sum_1 = op.ReduceSumAnyOpset(mul_18, init7_s1_1, keepdims=1, outputs=["sum_1"])
+    add_50 = op.Add(matmul, sum_1, outputs=["add_50"])
+    sum_2 = op.ReduceSumAnyOpset(mul_21, init7_s1_1, keepdims=1, outputs=["sum_2"])
+    permute_2 = op.Reshape(sum_2, init7_s2_1__1, outputs=["permute_2"])
+    add_59 = op.Add(add_50, permute_2, outputs=["add_59"])
+    _to_copy_1 = op.Cast(isnan_1, to=itype, outputs=["_to_copy_1"])
+    matmul_1 = op.Gemm(mul_18, _to_copy_1, transA=0, transB=1, outputs=["matmul_1"])
+    sub_32 = op.Sub(add_59, matmul_1, outputs=["sub_32"])
+    sub_43 = op.Sub(sub_32, matmul_2, outputs=["sub_43"])
+    clip = op.Clip(sub_43, init1_s1_, outputs=["clip"])
+    _reshape_init1_s_20 = op.Reshape(
+        init1_s_2, init7_s1_1, outputs=["_reshape_init1_s_20"]
+    )
+    rsub = op.Sub(_reshape_init1_s_20, _to_copy_2, outputs=["rsub"])
+    bitwise_not = op.Not(isnan_1, outputs=["bitwise_not"])
+    _to_copy_4 = op.Cast(bitwise_not, to=itype, outputs=["_to_copy_4"])
+    matmul_3 = op.Gemm(rsub, _to_copy_4, transA=0, transB=1, outputs=["matmul_3"])
+    _reshape_init1_s_30 = op.Reshape(
+        init1_s_3, init7_s1_1, outputs=["_reshape_init1_s_30"]
+    )
+    eq_61 = op.Equal(matmul_3, _reshape_init1_s_30, outputs=["eq_61"])
+    index_put_2 = op.Where(eq_61, c_lifted_tensor_2, clip, outputs=["index_put_2"])
+    maximum = op.Max(c_lifted_tensor_3, matmul_3, outputs=["maximum"])
+    div = op.Div(index_put_2, maximum, outputs=["div"])
+    col_size = op.Shape(x, start=1, end=2, outputs=["col_size"])
+    col_size_c = op.Cast(col_size, to=itype)
+    _onx_mul_div0 = op.Mul(div, col_size_c, outputs=["_onx_mul_div0"])
+    output_0 = op.Sqrt(_onx_mul_div0, outputs=["output_0"])
+    gr.make_tensor_output(output_0)
+    g.make_local_function(
+        container=gr, optimize=False, name="dist_nan_euclidean", domain="local_domain"
+    )
+    return gr
+
+
+def make_calc_impute_donors(g: ModelComponentContainer, scope: Scope, itype: int):
+    gr = ModelComponentContainer({"": g.main_opset}, as_function=True)
+    dist_pot_donors = gr.make_tensor_input("dist_pot_donors")
+    n_neighbors = gr.make_tensor_input("n_neighbors")
+    op = gr.get_op_builder(scope)
+    init7_s_0 = op.Constant(
+        value=from_array(np.array(0, dtype=np.int64), name="value"),
+        outputs=["init7_s_0"],
+    )
+    init7_s1_1 = op.Constant(
+        value=from_array(np.array([1], dtype=np.int64), name="value"),
+        outputs=["init7_s1_1"],
+    )
+    init7_s_1 = op.Constant(
+        value=from_array(np.array(1, dtype=np.int64), name="value"),
+        outputs=["init7_s_1"],
+    )
+    _shape_dist_pot_donors0 = op.Shape(
+        dist_pot_donors, end=1, start=0, outputs=["_shape_dist_pot_donors0"]
+    )
+    sym_size_int_4 = op.SqueezeAnyOpset(
+        _shape_dist_pot_donors0, outputs=["sym_size_int_4"]
+    )
+    if g.target_opset < 11:
+        unused_topk_values, neg_output_0 = op.TopK(
+            op.Neg(dist_pot_donors),
+            op.Reshape(n_neighbors, np.array([1], dtype=np.int64)),
+            largest=1,
+            sorted=1,
+            outputs=["unused_topk_values", "neg_output_0"],
+        )
+        output_0 = op.Neg(neg_output_0, outputs=["output_0"])
+    else:
+        max_value = (
+            np.finfo(np.float32).max
+            if itype == TensorProto.FLOAT
+            else np.finfo(np.float64).max
+        )
+        init_max_value = op.Constant(
+            value=from_array(
+                np.array(max_value, dtype=tensor_dtype_to_np_dtype(itype)), name="value"
+            ),
+            outputs=["init_max_value"],
+        )
+
+        unused_topk_values, output_0 = op.TopK(
+            op.Where(op.IsNaN(dist_pot_donors), init_max_value, dist_pot_donors),
+            op.Reshape(n_neighbors, np.array([1], dtype=np.int64)),
+            largest=0,
+            sorted=1,
+            outputs=["unused_topk_values", "output_0"],
+        )
+
+    arange = op.Range(init7_s_0, sym_size_int_4, init7_s_1, outputs=["arange"])
+    unsqueeze = op.UnsqueezeAnyOpset(arange, init7_s1_1, outputs=["unsqueeze"])
+    _onx_gathernd_dist_pot_donors0 = op.GatherND(
+        dist_pot_donors,
+        unsqueeze,
+        batch_dims=0,
+        outputs=["_onx_gathernd_dist_pot_donors0"],
+    )
+    output_1 = op.GatherElements(
+        _onx_gathernd_dist_pot_donors0, output_0, axis=1, outputs=["output_1"]
+    )
+    gr.make_tensor_output(output_0)
+    gr.make_tensor_output(output_1)
+    g.make_local_function(
+        container=gr, optimize=False, name="calc_impute_donors", domain="local_domain"
+    )
+    return gr
+
+
+def make_calc_impute_weights(g: ModelComponentContainer, scope: Scope, itype: int):
+    gr = ModelComponentContainer({"": g.main_opset}, as_function=True)
+    donors_dist = gr.make_tensor_input("donors_dist")
+    op = gr.get_op_builder(scope)
+    c_lifted_tensor_0 = op.Constant(
+        value=from_array(
+            np.array(0.0, dtype=tensor_dtype_to_np_dtype(itype)), name="value"
+        ),
+        outputs=["c_lifted_tensor_0"],
+    )
+    _shape_donors_dist0 = op.Shape(donors_dist, outputs=["_shape_donors_dist0"])
+    ones_like = op.ConstantOfShape(
+        _shape_donors_dist0,
+        value=from_array(
+            np.array([1.0], dtype=tensor_dtype_to_np_dtype(itype)), name="value"
+        ),
+        outputs=["ones_like"],
+    )
+    isnan = op.IsNaN(donors_dist, outputs=["isnan"])
+    output_0 = op.Where(isnan, c_lifted_tensor_0, ones_like, outputs=["output_0"])
+    gr.make_tensor_output(output_0)
+    g.make_local_function(
+        container=gr, optimize=False, name="calc_impute_weights", domain="local_domain"
+    )
+    return gr
+
+
+def make_calc_impute_make_new_neights(
+    g: ModelComponentContainer, scope: Scope, itype: int
+):
+    gr = ModelComponentContainer({"": g.main_opset}, as_function=True)
+    donors_mask = gr.make_tensor_input("donors_mask")
+    _donors = gr.make_tensor_input("donors")
+    weight_matrix = gr.make_tensor_input("weight_matrix")
+    op = gr.get_op_builder(scope)
+    _to_copy = op.Cast(donors_mask, to=itype, outputs=["_to_copy"])
+    output_0 = op.Mul(_to_copy, weight_matrix, outputs=["output_0"])
+    gr.make_tensor_output(output_0)
+    g.make_local_function(
+        container=gr,
+        optimize=False,
+        name="calc_impute_make_new_neights",
+        domain="local_domain",
+    )
+    return gr
+
+
+def make_calc_impute(g: ModelComponentContainer, scope: Scope, itype: int):
+    gr = ModelComponentContainer(
+        {"": g.main_opset, "local_domain": 1}, as_function=True
+    )
+    dist_pot_donors = gr.make_tensor_input("dist_pot_donors")
+    n_neighbors = gr.make_tensor_input("n_neighbors")
+    fit_x_col = gr.make_tensor_input("fit_x_col")
+    mask_fit_x_col = gr.make_tensor_input("mask_fit_x_col")
+    op = gr.get_op_builder(scope)
+    c_lifted_tensor_0 = op.Constant(
+        value=from_array(np.array([1], dtype=np.int64), name="value"),
+        outputs=["c_lifted_tensor_0"],
+    )
+    c_lifted_tensor_1 = op.Constant(
+        value=from_array(
+            np.array([1.0], dtype=tensor_dtype_to_np_dtype(itype)), name="value"
+        ),
+        outputs=["c_lifted_tensor_1"],
+    )
+    init7_s1__1 = op.Constant(
+        value=from_array(np.array([-1], dtype=np.int64), name="value"),
+        outputs=["init7_s1__1"],
+    )
+    init1_s_ = op.Constant(
+        value=from_array(
+            np.array(0.0, dtype=tensor_dtype_to_np_dtype(itype)), name="value"
+        ),
+        outputs=["init1_s_"],
+    )
+    (
+        calc_impute_in__donors_idx__0,
+        calc_impute_in__donors_idx__1,
+    ) = op.calc_impute_donors(
+        dist_pot_donors,
+        n_neighbors,
+        domain="local_domain",
+        outputs=[
+            "calc_impute_in__donors_idx__0",
+            "calc_impute_in__donors_idx__1",
+        ],
+    )
+    calc_impute_in__weights = op.calc_impute_weights(
+        calc_impute_in__donors_idx__1,
+        domain="local_domain",
+        outputs=["calc_impute_in__weights"],
+    )
+    _reshape_fit_x_col0 = op.Reshape(
+        fit_x_col, init7_s1__1, outputs=["_reshape_fit_x_col0"]
+    )
+    take = op.Gather(
+        _reshape_fit_x_col0,
+        calc_impute_in__donors_idx__0,
+        outputs=["take"],
+    )
+    _reshape_mask_fit_x_col0 = op.Reshape(
+        mask_fit_x_col, init7_s1__1, outputs=["_reshape_mask_fit_x_col0"]
+    )
+    take_1 = op.Gather(
+        _reshape_mask_fit_x_col0,
+        calc_impute_in__donors_idx__0,
+        outputs=["take_1"],
+    )
+    _to_copy = op.Cast(take_1, to=TensorProto.INT64, outputs=["_to_copy"])
+    sub_12 = op.Sub(c_lifted_tensor_0, _to_copy, outputs=["sub_12"])
+    calc_impute_in__make_new_neights = op.calc_impute_make_new_neights(
+        sub_12,
+        take,
+        calc_impute_in__weights,
+        domain="local_domain",
+        outputs=["calc_impute_in__make_new_neights"],
+    )
+    sum_1 = op.ReduceSumAnyOpset(
+        calc_impute_in__make_new_neights,
+        c_lifted_tensor_0,
+        keepdims=1,
+        outputs=["sum_1"],
+    )
+    _reshape_init1_s_0 = op.Reshape(
+        init1_s_, c_lifted_tensor_0, outputs=["_reshape_init1_s_0"]
+    )
+    eq_17 = op.Equal(sum_1, _reshape_init1_s_0, outputs=["eq_17"])
+    where = op.Where(eq_17, c_lifted_tensor_1, sum_1, outputs=["where"])
+    mul_17 = op.Mul(
+        take,
+        calc_impute_in__make_new_neights,
+        outputs=["mul_17"],
+    )
+    sum_2 = op.ReduceSumAnyOpset(
+        mul_17, c_lifted_tensor_0, keepdims=1, outputs=["sum_2"]
+    )
+    div = op.Div(sum_2, where, outputs=["div"])
+    output_0 = op.SqueezeAnyOpset(div, c_lifted_tensor_0, outputs=["output_0"])
+    gr.make_tensor_output(output_0)
+    g.make_local_function(
+        container=gr, optimize=False, name="calc_impute", domain="local_domain"
+    )
+    return gr
+
+
+def make_knn_imputer_column_all_nan(
+    g: ModelComponentContainer, scope: Scope, itype: int
+):
+    """
+    Sent:
+
+    ::
+
+        i_col,
+        x,
+        dist_subset,
+        mask_fit_x,
+        _fit_x,
+        receivers_idx,
+        all_nan_receivers_idx,
+        all_nan_dist_mask,
+        dist_chunk,
+        dist_idx_map,
+        potential_donors_idx,
+    """
+    gr = ModelComponentContainer(
+        {"": g.main_opset, "local_domain": 1}, as_function=True
+    )
+    i_col = gr.make_tensor_input("i_col")
+    x = gr.make_tensor_input("x")
+    _dist_subset = gr.make_tensor_input("dist_subset")
+    mask_fit_x = gr.make_tensor_input("mask_fit_x")
+    _fit_x = gr.make_tensor_input("_fit_x")
+    receivers_idx = gr.make_tensor_input("receivers_idx")
+    all_nan_receivers_idx = gr.make_tensor_input("all_nan_receivers_idx")
+    all_nan_dist_mask = gr.make_tensor_input("all_nan_dist_mask")
+    dist_chunk = gr.make_tensor_input("dist_chunk")
+    dist_idx_map = gr.make_tensor_input("dist_idx_map")
+    potential_donors_idx = gr.make_tensor_input("potential_donors_idx")
+    op = gr.get_op_builder(scope)
+    c_lifted_tensor_0 = op.Constant(
+        value=from_array(np.array([1.0], dtype=np.float32), name="value"),
+        outputs=["c_lifted_tensor_0"],
+    )
+    # init7_s_2 = op.Constant(value=from_array(
+    # np.array(2, dtype=np.int64), name='value'), outputs=['init7_s_2'])
+    init7_s_2 = i_col
+    init1_s_ = op.Constant(
+        value=from_array(np.array(1.0, dtype=np.float32), name="value"),
+        outputs=["init1_s_"],
+    )
+    init7_s1_1 = op.Constant(
+        value=from_array(np.array([1], dtype=np.int64), name="value"),
+        outputs=["init7_s1_1"],
+    )
+    init1_s_2 = op.Constant(
+        value=from_array(np.array(0.0, dtype=np.float32), name="value"),
+        outputs=["init1_s_2"],
+    )
+    init7_s1_0 = op.Constant(
+        value=from_array(np.array([0], dtype=np.int64), name="value"),
+        outputs=["init7_s1_0"],
+    )
+    init7_s1__1 = op.Constant(
+        value=from_array(np.array([-1], dtype=np.int64), name="value"),
+        outputs=["init7_s1__1"],
+    )
+    init7_s_1 = op.Constant(
+        value=from_array(np.array(1, dtype=np.int64), name="value"),
+        outputs=["init7_s_1"],
+    )
+    # init7_s1_2 = op.Constant(value=from_array(np.array([2], dtype=np.int64),
+    # name='value'), outputs=['init7_s1_2'])
+    init7_s1_2 = op.UnsqueezeAnyOpset(i_col, init7_s1_0, outputs=["init7_s1_2"])
+    select = op.Gather(mask_fit_x, init7_s_2, axis=1, outputs=["select"])
+    bitwise_not = op.Not(select, outputs=["bitwise_not"])
+    _to_copy = op.Cast(bitwise_not, to=1, outputs=["_to_copy"])
+    sum_1 = op.ReduceSumAnyOpset(_to_copy, keepdims=0, outputs=["sum_1"])
+    _reshape_init1_s_0 = op.Reshape(
+        init1_s_, init7_s1_1, outputs=["_reshape_init1_s_0"]
+    )
+    eq_6 = op.Equal(_to_copy, _reshape_init1_s_0, outputs=["eq_6"])
+    select_1 = op.Gather(_fit_x, init7_s_2, axis=1, outputs=["select_1"])
+    index = op.Compress(select_1, eq_6, axis=0, outputs=["index"])
+    sum_2 = op.ReduceSumAnyOpset(index, keepdims=0, outputs=["sum_2"])
+    gt = op.Greater(sum_1, init1_s_2, outputs=["gt"])
+    where = op.Where(gt, sum_1, c_lifted_tensor_0, outputs=["where"])
+    _reshape__to_copy_20 = op.Reshape(
+        sum_2, init7_s1_1, outputs=["_reshape__to_copy_20"]
+    )
+    div = op.Div(_reshape__to_copy_20, where, outputs=["div"])
+    select_2 = op.Gather(x, init7_s_2, axis=1, outputs=["select_2"])
+    view = op.SqueezeAnyOpset(div, init7_s1_0, outputs=["view"])
+    _onx_unsqueeze_all_nan_receivers_idx0 = op.UnsqueezeAnyOpset(
+        all_nan_receivers_idx,
+        init7_s1__1,
+        outputs=["_onx_unsqueeze_all_nan_receivers_idx0"],
+    )
+    _shape_all_nan_receivers_idx0 = op.Shape(
+        all_nan_receivers_idx, outputs=["_shape_all_nan_receivers_idx0"]
+    )
+    _onx_expand_view0 = op.Expand(
+        view, _shape_all_nan_receivers_idx0, outputs=["_onx_expand_view0"]
+    )
+    index_put = op.ScatterND(
+        select_2,
+        _onx_unsqueeze_all_nan_receivers_idx0,
+        _onx_expand_view0,
+        outputs=["index_put"],
+    )
+    _onx_unsqueeze_index_put0 = op.UnsqueezeAnyOpset(
+        index_put, init7_s_1, outputs=["_onx_unsqueeze_index_put0"]
+    )
+    _shape_unsqueeze_index_put00 = op.Shape(
+        _onx_unsqueeze_index_put0, outputs=["_shape_unsqueeze_index_put00"]
+    )
+    _onx_expand_init7_s1_20 = op.Expand(
+        init7_s1_2, _shape_unsqueeze_index_put00, outputs=["_onx_expand_init7_s1_20"]
+    )
+    output_0 = op.ScatterElements(
+        x,
+        _onx_expand_init7_s1_20,
+        _onx_unsqueeze_index_put0,
+        axis=1,
+        reduction="none",
+        outputs=["output_0"],
+    )
+    bitwise_not_1 = op.Not(all_nan_dist_mask, outputs=["bitwise_not_1"])
+    output_2 = op.Compress(receivers_idx, bitwise_not_1, axis=0, outputs=["output_2"])
+    index_2 = op.Gather(dist_idx_map, output_2, axis=0, outputs=["index_2"])
+    index_3 = op.Gather(dist_chunk, index_2, axis=0, outputs=["index_3"])
+    output_1 = op.Gather(index_3, potential_donors_idx, axis=1, outputs=["output_1"])
+    gr.make_tensor_output(output_0)
+    gr.make_tensor_output(output_1)
+    gr.make_tensor_output(output_2)
+
+    g.make_local_function(
+        container=gr,
+        optimize=False,
+        name="knn_imputer_column_all_nan",
+        domain="local_domain",
+    )
+    return gr
+
+
+def make_knn_imputer_column_nan_found(
+    g: ModelComponentContainer, scope: Scope, itype: int
+):
+    gr = ModelComponentContainer(
+        {"": g.main_opset, "local_domain": 1}, as_function=True
+    )
+    i_col = gr.make_tensor_input("i_col")
+    x = gr.make_tensor_input("x")
+    col_mask = gr.make_tensor_input("col_mask")
+    dist_idx_map = gr.make_tensor_input("dist_idx_map")
+    mask_fit_x = gr.make_tensor_input("mask_fit_x")
+    non_missing_fix_x = gr.make_tensor_input("non_missing_fix_x")
+    dist_chunk = gr.make_tensor_input("dist_chunk")
+    _fit_x = gr.make_tensor_input("_fit_x")
+    row_missing_idx = gr.make_tensor_input("row_missing_idx")
+
+    op = gr.get_op_builder(scope)
+
+    init7_s_2 = i_col
+    init7_s1__1 = op.Constant(
+        value=from_array(np.array([-1], dtype=np.int64), name="value"),
+        outputs=["init7_s1__1"],
+    )
+    init7_s_0 = op.Constant(
+        value=from_array(np.array(0, dtype=np.int64), name="value"),
+        outputs=["init7_s_0"],
+    )
+    init7_s_1 = op.Constant(
+        value=from_array(np.array(1, dtype=np.int64), name="value"),
+        outputs=["init7_s_1"],
+    )
+    c_lifted_tensor_1 = op.Constant(
+        value=from_array(np.array(3, dtype=np.int64), name="value"),
+        outputs=["c_lifted_tensor_1"],
+    )
+    c_lifted_tensor_2 = op.Constant(
+        value=from_array(np.array([1], dtype=np.int64), name="value"),
+        outputs=["c_lifted_tensor_2"],
+    )
+
+    init7_s1_2 = op.UnsqueezeAnyOpset(init7_s_2, np.array([0], dtype=np.int64))
+
+    select_1 = op.Gather(non_missing_fix_x, init7_s_2, axis=1, outputs=["select_1"])
+    potential_donors_idx = op.NonZero(select_1, outputs=["potential_donors_idx"])
+
+    nonzero_numpy__0 = op.Reshape(
+        potential_donors_idx, init7_s1__1, outputs=["nonzero_numpy__0"]
+    )
+    _shape_getitem_20 = op.Shape(
+        nonzero_numpy__0, end=1, start=0, outputs=["_shape_getitem_20"]
+    )
+    sym_size_int_23 = op.SqueezeAnyOpset(_shape_getitem_20, outputs=["sym_size_int_23"])
+
+    flat_index = op.NonZero(col_mask, outputs=["flat_index"])
+    nonzero_numpy_1__0 = op.Reshape(
+        flat_index, init7_s1__1, outputs=["nonzero_numpy_1__0"]
+    )
+    receivers_idx = op.Gather(
+        row_missing_idx, nonzero_numpy_1__0, axis=0, outputs=["receivers_idx"]
+    )
+    index_2 = op.Gather(dist_idx_map, receivers_idx, axis=0, outputs=["index_2"])
+    index_3 = op.Gather(dist_chunk, index_2, axis=0, outputs=["index_3"])
+    dist_subset = op.Gather(index_3, nonzero_numpy__0, axis=1, outputs=["dist_subset"])
+    isnan = op.IsNaN(dist_subset, outputs=["isnan"])
+    _onx_cast_isnan0 = op.Cast(isnan, to=6, outputs=["_onx_cast_isnan0"])
+    _onx_reducemin_cast_isnan00 = op.ReduceMinAnyOpset(
+        _onx_cast_isnan0,
+        c_lifted_tensor_2,
+        keepdims=0,
+        outputs=["_onx_reducemin_cast_isnan00"],
+    )
+    all_nan_dist_mask = op.Cast(
+        _onx_reducemin_cast_isnan00, to=TensorProto.BOOL, outputs=["all_nan_dist_mask"]
+    )
+    all_nan_receivers_idx = op.Compress(
+        receivers_idx, all_nan_dist_mask, axis=0, outputs=["all_nan_receivers_idx"]
+    )
+
+    size_all_nan_receivers_idx = op.Size(all_nan_receivers_idx)
+    zero_i = op.Constant(value=from_array(np.array(0, dtype=np.int64), name="zero"))
+    size_index_5_greater = op.Greater(size_all_nan_receivers_idx, zero_i)
+    select_scatter, dist_subset, receivers_idx = op.If(
+        size_index_5_greater,
+        then_branch=make_graph(
+            [
+                make_node(
+                    "knn_imputer_column_all_nan",
+                    [
+                        i_col,
+                        x,
+                        dist_subset,
+                        mask_fit_x,
+                        _fit_x,
+                        receivers_idx,
+                        all_nan_receivers_idx,
+                        all_nan_dist_mask,
+                        dist_chunk,
+                        dist_idx_map,
+                        potential_donors_idx,
+                    ],
+                    ["A", "B", "C"],
+                    domain="local_domain",
+                )
+            ],
+            "then_branch",
+            [],
+            [
+                make_tensor_value_info("A", itype, None),
+                make_tensor_value_info("B", itype, None),
+                make_tensor_value_info("C", TensorProto.INT64, None),
+            ],
+        ),
+        else_branch=make_graph(
+            [
+                make_node("Identity", [x], ["A"]),
+                make_node("Identity", [dist_subset], ["B"]),
+                make_node("Identity", [receivers_idx], ["C"]),
+            ],
+            "identity",
+            [],
+            [
+                make_tensor_value_info("A", itype, None),
+                make_tensor_value_info("B", itype, None),
+                make_tensor_value_info("C", TensorProto.INT64, None),
+            ],
+        ),
+        outputs=["select_scatter", "dist_subset_updated", "receivers_idx_updated"],
+    )
+
+    lt = op.Less(c_lifted_tensor_1, sym_size_int_23, outputs=["lt"])
+    where_1 = op.Where(lt, c_lifted_tensor_1, sym_size_int_23, outputs=["where_1"])
+    le = op.LessOrEqual(where_1, init7_s_0, outputs=["le"])
+    # c_lifted_tensor_2 -> init7_s_1 to have a zero time, onnxruntime crashes wher shapes are
+    # (), (1,), ()
+    where_2 = op.Where(le, init7_s_1, where_1, outputs=["where_2"])
+    select_6 = op.Gather(_fit_x, init7_s_2, axis=1, outputs=["select_6"])
+    index_11 = op.Gather(select_6, nonzero_numpy__0, axis=0, outputs=["index_11"])
+    select_7 = op.Gather(mask_fit_x, init7_s_2, axis=1, outputs=["select_7"])
+    index_12 = op.Gather(select_7, nonzero_numpy__0, axis=0, outputs=["index_12"])
+    calc_impute_output = op.calc_impute(
+        dist_subset,
+        where_2,
+        index_11,
+        index_12,
+        domain="local_domain",
+        outputs=["calc_impute_output"],
+    )
+    select_9 = op.Gather(select_scatter, init7_s_2, axis=1, outputs=["select_9"])
+    _onx_unsqueeze_index_70 = op.UnsqueezeAnyOpset(
+        receivers_idx, init7_s1__1, outputs=["_onx_unsqueeze_index_70"]
+    )
+    index_put_1 = op.ScatterND(
+        select_9,
+        _onx_unsqueeze_index_70,
+        calc_impute_output,
+        outputs=["index_put_1"],
+    )
+    _onx_unsqueeze_index_put_10 = op.UnsqueezeAnyOpset(
+        index_put_1, init7_s_1, outputs=["_onx_unsqueeze_index_put_10"]
+    )
+    _shape_unsqueeze_index_put_100 = op.Shape(
+        _onx_unsqueeze_index_put_10, outputs=["_shape_unsqueeze_index_put_100"]
+    )
+    _onx_expand_init7_s1_202 = op.Expand(
+        init7_s1_2, _shape_unsqueeze_index_put_100, outputs=["_onx_expand_init7_s1_202"]
+    )
+    output_0 = op.ScatterElements(
+        select_scatter,
+        _onx_expand_init7_s1_202,
+        _onx_unsqueeze_index_put_10,
+        axis=1,
+        reduction="none",
+        outputs=["output_0"],
+    )
+    gr.make_tensor_output(output_0)
+    g.make_local_function(
+        container=gr,
+        optimize=False,
+        name="knn_imputer_column_nan_found",
+        domain="local_domain",
+    )
+    return gr
+
+
+def make_knn_imputer_column(g: ModelComponentContainer, scope: Scope, itype: int):
+    gr = ModelComponentContainer(
+        {"": g.main_opset, "local_domain": 1}, as_function=True
+    )
+    i_col = gr.make_tensor_input("i_col")
+    x = gr.make_tensor_input("x")
+    dist_chunk = gr.make_tensor_input("dist_chunk")
+    non_missing_fix_x = gr.make_tensor_input("non_missing_fix_x")
+    mask_fit_x = gr.make_tensor_input("mask_fit_x")
+    dist_idx_map = gr.make_tensor_input("dist_idx_map")
+    mask = gr.make_tensor_input("mask")
+    row_missing_idx = gr.make_tensor_input("row_missing_idx")
+    _fit_x = gr.make_tensor_input("_fit_x")
+
+    op = gr.get_op_builder(scope)
+    zero32 = op.Constant(
+        value=from_array(np.array(0, dtype=np.int32), name="value"),
+        outputs=["zero"],
+    )
+    init7_s_2 = i_col
+    # init7_s_2 = op.Constant(value=from_array(np.array(2, dtype=np.int64), name='value'),
+    # outputs=['init7_s_2'])
+    init7_s1__1 = op.Constant(
+        value=from_array(np.array([-1], dtype=np.int64), name="value"),
+        outputs=["init7_s1__1"],
+    )
+    # init7_s1_2 = op.Constant(value=from_array(np.array([2], dtype=np.int64), name='value'),
+    # outputs=['init7_s1_2'])
+    select = op.Gather(mask, init7_s_2, axis=1, outputs=["select"])
+    index = op.Gather(select, row_missing_idx, axis=0, outputs=["index"])
+    view = op.Reshape(index, init7_s1__1, outputs=["view"])
+    col_mask = op.Cast(view, to=TensorProto.INT32)
+
+    # If max(col_mask) == 0, no need to continue
+    # That avoids dealing with empty sizes.
+    col_mask_max = op.ReduceMax(col_mask, outputs=["col_mask_max"])
+    col_mask_max_null = op.Equal(col_mask_max, zero32, outputs=["col_mask_max_null"])
+
+    output_0 = op.If(
+        col_mask_max_null,
+        then_branch=make_graph(
+            [make_node("Identity", [x], ["X"])],
+            "identity",
+            [],
+            [make_tensor_value_info("X", itype, None)],
+        ),
+        else_branch=make_graph(
+            [
+                make_node(
+                    "knn_imputer_column_nan_found",
+                    [
+                        i_col,
+                        x,
+                        col_mask,
+                        dist_idx_map,
+                        mask_fit_x,
+                        non_missing_fix_x,
+                        dist_chunk,
+                        _fit_x,
+                        row_missing_idx,
+                    ],
+                    ["X"],
+                    domain="local_domain",
+                )
+            ],
+            "then_branch",
+            [],
+            [make_tensor_value_info("X", itype, None)],
+        ),
+        outputs=["output_0"],
+    )
+    gr.make_tensor_output(output_0)
+    g.make_local_function(
+        container=gr, optimize=False, name="knn_imputer_column", domain="local_domain"
+    )
+    return gr
+
+
+def _knn_imputer_builder(
+    op: ModelComponentContainer,
+    _mask_fit_x: "BOOL[s0, 2]",  # noqa: F821
+    _valid_mask: "BOOL[2]",  # noqa: F821
+    _fit_x: "DOUBLE[s1, 2]",  # noqa: F821
+    x: "FLOAT[s2, 2]",  # noqa: F821
+    itype: int,
+):
+    n_cols = _fit_x.shape[1]
+    init7_s1_1 = np.array([1], dtype=np.int64)
+    init7_s1__1 = np.array([-1], dtype=np.int64)
+    isnan = op.IsNaN(x, outputs=["isnan"])
+    _onx_compress_isnan0 = op.Compress(
+        isnan, _valid_mask, axis=1, outputs=["_onx_compress_isnan0"]
+    )
+    _onx_cast_index0 = op.Cast(_onx_compress_isnan0, to=6, outputs=["_onx_cast_index0"])
+    _onx_reducemax_cast_index00 = op.ReduceMaxAnyOpset(
+        _onx_cast_index0,
+        init7_s1_1,
+        keepdims=0,
+        outputs=["_onx_reducemax_cast_index00"],
+    )
+    any_1 = op.Cast(_onx_reducemax_cast_index00, to=TensorProto.BOOL, outputs=["any_1"])
+    view = op.Reshape(any_1, init7_s1__1, outputs=["view"])
+    _onx_nonzero_view0 = op.NonZero(view, outputs=["_onx_nonzero_view0"])
+    nonzero_numpy__0 = op.Reshape(
+        _onx_nonzero_view0, init7_s1__1, outputs=["nonzero_numpy__0"]
+    )
+    logical_not = op.Not(_mask_fit_x, outputs=["logical_not"])
+    make_dict_idx_map = op.dict_idx_map(
+        x, nonzero_numpy__0, domain="local_domain", outputs=["make_dict_idx_map"]
+    )
+    index_1 = op.Gather(x, nonzero_numpy__0, axis=0, outputs=["index_1"])
+    c_torch_knnimputer_dist = op.dist_nan_euclidean(
+        index_1, _fit_x, domain="local_domain", outputs=["c_torch_knnimputer_dist"]
+    )
+    for i_cols in range(n_cols):
+        x = op.knn_imputer_column(
+            np.array(i_cols, dtype=np.int64),
+            x,
+            c_torch_knnimputer_dist,
+            logical_not,
+            _mask_fit_x,
+            make_dict_idx_map,
+            isnan,
+            nonzero_numpy__0,
+            _fit_x,
+            domain="local_domain",
+        )
+    return op.Compress(x, _valid_mask, axis=1, outputs=["output_0"])
+
+
 def convert_knn_imputer(
     scope: Scope, operator: Operator, container: ModelComponentContainer
 ):
     """
     Converts *KNNImputer* into *ONNX*.
     """
+    assert (
+        container.target_opset >= 18
+    ), f"This converter no longer works for opset {container.target_opset} < 18."
     dtype = guess_numpy_type(operator.inputs[0].type)
     if dtype != np.float64:
         dtype = np.float32
@@ -804,63 +1638,37 @@ def convert_knn_imputer(
     if knn_op.metric != "nan_euclidean":
         raise RuntimeError("Unable to convert KNNImputer when metric is callable.")
     if knn_op.weights not in ("uniform", "distance"):
-        raise RuntimeError("Unable to convert KNNImputer when weights is callable.")
+        raise RuntimeError(
+            f"Unable to convert KNNImputer when weights "
+            f"is callable, knn_op.weights={knn_op.weights}"
+        )
     if knn_op.weights == "distance":
         raise NotImplementedError(
             "KNNImputer with distance as metric is not supported, "
             "you may raise an issue at "
             "https://github.com/onnx/sklearn-onnx/issues."
         )
-    options = container.get_options(knn_op, dict(optim=None))
-    op_version = container.target_opset
-    input_name = operator.inputs[0]
+    # options = container.get_options(knn_op, dict(optim=None))
+    # options are not used anymore
     training_data = knn_op._fit_X.astype(dtype)
-    training_data[np.isnan(training_data)] = 0
-    out = operator.outputs
-    top_indices, missing_input_name = _nearest_neighbours(
-        container, knn_op, input_name, op_version, options["optim"], dtype, proto_type
-    )
-    flattened = OnnxFlatten(top_indices, op_version=op_version)
-    extracted = OnnxArrayFeatureExtractor(
-        training_data.T, flattened, op_version=op_version
-    )
-    reshaped = OnnxReshapeApi13(
-        extracted,
-        np.array([training_data.shape[1], -1, knn_op.n_neighbors], dtype=np.int64),
-        op_version=op_version,
-    )
-    transpose_result = OnnxTranspose(reshaped, op_version=op_version, perm=[1, 2, 0])
-    reduced = OnnxReduceSumApi11(
-        transpose_result, op_version=op_version, axes=[1], keepdims=0
-    )
-    cast_res = OnnxCast(
-        OnnxCast(
-            transpose_result, to=onnx_proto.TensorProto.BOOL, op_version=op_version
-        ),
-        to=proto_type,
-        op_version=op_version,
-    )
-    deno = OnnxReduceSumApi11(cast_res, op_version=op_version, axes=[1], keepdims=0)
-    deno_updated = OnnxAdd(
-        deno,
-        OnnxCast(
-            OnnxNot(
-                OnnxCast(deno, to=onnx_proto.TensorProto.BOOL, op_version=op_version),
-                op_version=op_version,
-            ),
-            to=proto_type,
-            op_version=op_version,
-        ),
-        op_version=op_version,
-    )
-    imputed_out = OnnxWhere(
-        missing_input_name,
-        OnnxDiv(reduced, deno_updated, op_version=op_version),
-        input_name,
-        output_names=out[:1],
-        op_version=op_version,
+    result = _knn_imputer_builder(
+        container.get_op_builder(scope),
+        knn_op._mask_fit_X,
+        knn_op._valid_mask,
+        training_data,
+        operator.inputs[0].full_name,
+        itype=proto_type,
     )
-    imputed_out.add_to(scope, container)
+    make_dict_idx_map(container, Scope("LF1"), itype=proto_type)
+    make_dist_nan_euclidean(container, Scope("LF2"), itype=proto_type)
+    make_calc_impute_donors(container, Scope("LF3"), itype=proto_type)
+    make_calc_impute_weights(container, Scope("LF4"), itype=proto_type)
+    make_calc_impute_make_new_neights(container, Scope("LF5"), itype=proto_type)
+    make_calc_impute(container, Scope("LF6"), itype=proto_type)
+    make_knn_imputer_column_all_nan(container, Scope("LF7"), itype=proto_type)
+    make_knn_imputer_column_nan_found(container, Scope("LF8"), itype=proto_type)
+    make_knn_imputer_column(container, Scope("LF9"), itype=proto_type)
+    container.add_node("Identity", [result], [operator.outputs[0].full_name])
 
 
 def convert_nca(scope: Scope, operator: Operator, container: ModelComponentContainer):
diff --git a/skl2onnx/proto/__init__.py b/skl2onnx/proto/__init__.py
index ab89e3e..6002812 100644
--- a/skl2onnx/proto/__init__.py
+++ b/skl2onnx/proto/__init__.py
@@ -19,7 +19,10 @@ except ImportError:
     # onnx is too old.
     pass
 
-from onnx.helper import split_complex_to_pairs
+try:
+    from onnx.helper import _split_complex_to_pairs as split_complex_to_pairs
+except ImportError:
+    from onnx.helper import split_complex_to_pairs
 
 
 def make_tensor_fixed(name, data_type, dims, vals, raw=False):


Test patch:
diff --git a/tests/test_convert.py b/tests/test_convert.py
index f672ebf..a9fa45e 100644
--- a/tests/test_convert.py
+++ b/tests/test_convert.py
@@ -47,7 +47,7 @@ class TestConvert(unittest.TestCase):
                     model, X[:1].astype(numpy.float32), target_opset=tops
                 )
                 dom = get_domain_opset(model_onnx)
-                self.assertEqual(len(dom), 1)
+                self.assertEqual(len(dom), 2)
                 assert dom[""] <= i
 
     @unittest.skipIf(KBinsDiscretizer is None, "skl too old")
@@ -82,7 +82,7 @@ class TestConvert(unittest.TestCase):
                     model, X[:1].astype(numpy.float32), target_opset=tops
                 )
                 dom = get_domain_opset(model_onnx)
-                self.assertEqual(len(dom), 1)
+                self.assertEqual(len(dom), 2)
                 self.assertIn(dom[""], (i, i - 1))
 
     def test_onehot(self):
diff --git a/tests/test_sklearn_nearest_neighbour_converter.py b/tests/test_sklearn_nearest_neighbour_converter.py
index d9ae8a8..27589ea 100644
--- a/tests/test_sklearn_nearest_neighbour_converter.py
+++ b/tests/test_sklearn_nearest_neighbour_converter.py
@@ -4,6 +4,8 @@
 Tests scikit-learn's KNeighbours Classifier and Regressor converters.
 """
 
+import math
+import numbers
 import sys
 import warnings
 import unittest
@@ -11,6 +13,7 @@ import functools
 import packaging.version as pv
 import numpy
 from numpy.testing import assert_almost_equal
+from onnx.reference import ReferenceEvaluator
 from onnxruntime import __version__ as ort_version
 from pandas import DataFrame
 
@@ -915,10 +918,484 @@ class TestNearestNeighbourConverter(unittest.TestCase):
             X_test, model, model_onnx, basename="SklearnKNNTransformerConnectivity"
         )
 
+    def _get_torch_knn_imputer(self):
+        import torch
+
+        def _get_weights(dist, weights):
+            """Get the weights from an array of distances and a parameter ``weights``.
+
+            Assume weights have already been validated.
+
+            Parameters
+            ----------
+            dist : ndarray
+                The input distances.
+
+            weights : {'uniform', 'distance'}, callable or None
+                The kind of weighting used.
+
+            Returns
+            -------
+            weights_arr : array of the same shape as ``dist``
+                If ``weights == 'uniform'``, then returns None.
+            """
+            if weights in (None, "uniform"):
+                return None
+
+            if weights == "distance":
+                # if user attempts to classify a point that was zero distance from one
+                # or more training points, those training points are weighted as 1.0
+                # and the other points as 0.0
+                dist = 1.0 / dist
+                inf_mask = torch.isinf(dist)
+                inf_row = torch.any(inf_mask, axis=1)
+                dist[inf_row] = inf_mask[inf_row]
+                return dist
+
+            if callable(weights):
+                return weights(dist)
+
+        def flatnonzero(x):
+            return torch.nonzero(torch.reshape(x, (-1,)), as_tuple=True)[0]
+
+        class NanEuclidean(torch.nn.Module):
+            """Implements :func:`sklearn.metrics.nan_euclidean_distances`."""
+
+            def __init__(self, squared=False, copy=True):
+                super().__init__()
+                self.squared = squared
+                self.copy = copy
+
+            def forward(self, X, Y):
+                X = X.clone()
+                Y = Y.to(X.dtype).clone()
+
+                missing_X = torch.isnan(X)
+                missing_Y = torch.isnan(Y)
+
+                # set missing values to zero
+                X[missing_X] = 0
+                Y[missing_Y] = 0
+
+                # Adjust distances for missing values
+                XX = X * X
+                YY = Y * Y
+
+                distances = -2 * X @ Y.T
+                distances += XX.sum(1, keepdim=True) + YY.sum(1, keepdim=True).T
+
+                distances -= XX @ missing_Y.to(X.dtype).T
+                distances -= missing_X.to(X.dtype) @ YY.T
+
+                distances = torch.clip(distances, 0, None)
+
+                present_X = 1 - missing_X.to(X.dtype)
+                present_Y = ~missing_Y
+                present_count = present_X @ present_Y.to(X.dtype).T
+                distances[present_count == 0] = torch.nan
+                # avoid divide by zero
+                present_count = torch.maximum(
+                    torch.tensor([1], dtype=present_count.dtype), present_count
+                )
+                distances /= present_count
+                distances *= X.shape[1]
+
+                if not self.squared:
+                    distances = distances.sqrt()
+
+                return distances
+
+        def _get_mask(X, value_to_mask):
+            return (
+                torch.isnan(X)
+                if (  # sklearn.utils._missing.is_scalar_nan(value_to_mask)
+                    not isinstance(value_to_mask, numbers.Integral)
+                    and isinstance(value_to_mask, numbers.Real)
+                    and math.isnan(value_to_mask)
+                )
+                else (value_to_mask == X)
+            )
+
+        class SubWeightMatrix(torch.nn.Module):
+            def __init__(self, weights):
+                super().__init__()
+                self.weights = weights
+
+            def forward(self, donors_dist):
+                weight_matrix = _get_weights(donors_dist, self.weights)
+                if weight_matrix is not None:
+                    weight_matrix = weight_matrix.clone()
+                    weight_matrix[torch.isnan(weight_matrix)] = 0.0
+                else:
+                    weight_matrix = torch.ones_like(donors_dist)
+                    weight_matrix[torch.isnan(donors_dist)] = 0.0
+                return weight_matrix
+
+        class SubDonorsIdx(torch.nn.Module):
+            def forward(self, dist_pot_donors, n_neighbors):
+                xn = torch.nan_to_num(dist_pot_donors, nan=1.0e10)
+                tk = torch.topk(xn, n_neighbors, dim=1, largest=False, sorted=True)
+                return tk.indices, tk.values
+
+        class MakeNewWeights(torch.nn.Module):
+            def forward(self, donors_mask, donors, weight_matrix):
+                return donors_mask.to(donors.dtype) * weight_matrix.to(donors.dtype)
+
+        class CalcImpute(torch.nn.Module):
+            """Implements :meth:`sklearn.impute.KNNImputer._calc_impute`."""
+
+            def __init__(self, weights):
+                super().__init__()
+                self._weights = SubWeightMatrix(weights)
+                self._donors_idx = SubDonorsIdx()
+                self._make_new_neights = MakeNewWeights()
+
+            def _calc_impute(
+                self, dist_pot_donors, n_neighbors, fit_X_col, mask_fit_X_col
+            ):
+                donors_idx, donors_dist = self._donors_idx(dist_pot_donors, n_neighbors)
+                weight_matrix = self._weights(donors_dist)
+                # Retrieve donor values and calculate kNN average
+                donors = fit_X_col.take(donors_idx)
+                donors_mask = torch.tensor([1], dtype=donors_idx.dtype) - (
+                    mask_fit_X_col.take(donors_idx)
+                ).to(donors_idx.dtype)
+
+                new_weights = self._make_new_neights(donors_mask, donors, weight_matrix)
+
+                weights_sum = new_weights.sum(axis=1, keepdim=True)
+                div = torch.where(
+                    weights_sum == 0,
+                    torch.tensor([1], dtype=weights_sum.dtype),
+                    weights_sum,
+                )
+                res = (donors * new_weights).sum(axis=1, keepdim=True) / div
+                return res.squeeze(dim=1).to(dist_pot_donors.dtype)
+
+            def forward(self, dist_pot_donors, n_neighbors, fit_X_col, mask_fit_X_col):
+                return self._calc_impute(
+                    dist_pot_donors, n_neighbors, fit_X_col, mask_fit_X_col
+                )
+
+        class ColProcessorAllNan(torch.nn.Module):
+            def __init__(self, col: int):
+                super().__init__()
+                self.col = col
+
+            def forward(
+                self,
+                X,
+                dist_subset,
+                mask_fit_X,
+                _fit_X,
+                receivers_idx,
+                all_nan_receivers_idx,
+                all_nan_dist_mask,
+                dist_chunk,
+                dist_idx_map,
+                potential_donors_idx,
+            ):
+                col = self.col
+                X = X.clone()
+                mask_ = (~mask_fit_X[:, col]).to(_fit_X.dtype)
+                mask_sum = mask_.to(X.dtype).sum()
+
+                col_sum = (_fit_X[mask_ == 1, col]).sum().to(X.dtype)
+                div = torch.where(
+                    mask_sum > 0, mask_sum, torch.tensor([1], dtype=mask_sum.dtype)
+                )
+                X[all_nan_receivers_idx, col] = col_sum / div
+
+                # receivers with at least one defined distance
+                receivers_idx = receivers_idx[~all_nan_dist_mask]
+                dist_subset = dist_chunk[dist_idx_map[receivers_idx]][
+                    :, potential_donors_idx
+                ]
+                return X, dist_subset, receivers_idx
+
+        class ColProcessorIdentity(torch.nn.Module):
+            def forward(
+                self,
+                X,
+                dist_subset,
+                mask_fit_X,
+                _fit_X,
+                receivers_idx,
+                all_nan_receivers_idx,
+                all_nan_dist_mask,
+                dist_chunk,
+                dist_idx_map,
+                potential_donors_idx,
+            ):
+                return (
+                    X,
+                    dist_subset,
+                    receivers_idx,
+                )
+
+        class ColProcessorCond(torch.nn.Module):
+            def __init__(self, col: int):
+                super().__init__()
+                self.col = col
+                self._all_nan = ColProcessorAllNan(col)
+                self._identity = ColProcessorIdentity()
+
+            def forward(
+                self,
+                X,
+                dist_subset,
+                mask_fit_X,
+                _fit_X,
+                receivers_idx,
+                all_nan_receivers_idx,
+                all_nan_dist_mask,
+                dist_chunk,
+                dist_idx_map,
+                potential_donors_idx,
+            ):
+                X, dist_subset, receivers_idx = torch.cond(
+                    all_nan_receivers_idx.numel() > 0,
+                    self._all_nan,
+                    self._identity,
+                    [
+                        X,
+                        dist_subset,
+                        mask_fit_X,
+                        _fit_X,
+                        receivers_idx,
+                        all_nan_receivers_idx,
+                        all_nan_dist_mask,
+                        dist_chunk,
+                        dist_idx_map,
+                        potential_donors_idx,
+                    ],
+                )
+                return X, dist_subset, receivers_idx
+
+        class ColProcessor(torch.nn.Module):
+            """Processes one column (= one feature)."""
+
+            def __init__(self, col, n_neighbors, weights):
+                super().__init__()
+                self._calc_impute = CalcImpute(weights)
+                self._col_cond = ColProcessorCond(col)
+                self.col = col
+                self.n_neighbors = n_neighbors
+
+            def process_one_col(
+                self,
+                X,
+                dist_chunk,
+                non_missing_fix_X,
+                mask_fit_X,
+                dist_idx_map,
+                mask,
+                row_missing_idx,
+                _fit_X,
+            ):
+                col = self.col
+                X = X.clone()
+                row_missing_chunk = row_missing_idx
+                col_mask = mask[row_missing_chunk, col]
+
+                potential_donors_idx = torch.nonzero(
+                    non_missing_fix_X[:, col], as_tuple=True
+                )[0]
+
+                # receivers_idx are indices in X
+                # if col_mask is all False, no need to continue
+                flat_index = flatnonzero(col_mask)
+
+                receivers_idx = row_missing_chunk[flat_index]
+
+                # distances for samples that needed imputation for column
+                dist_subset = dist_chunk[dist_idx_map[receivers_idx]][
+                    :, potential_donors_idx
+                ]
+
+                # receivers with all nan distances impute with mean
+                all_nan_dist_mask = torch.isnan(dist_subset).all(axis=1)
+                all_nan_receivers_idx = receivers_idx[all_nan_dist_mask]
+
+                # when all_nan_receivers_idx is not empty (training set is small)
+                # ... if all_nan_receivers_idx.size > 0:
+                #    # onnxruntime does not like this part when it is empty
+                #    mask_ = (~mask_fit_X[:, col]).to(_fit_X.dtype)
+                #    mask_sum = mask_.to(X.dtype).sum()
+                #
+                #    col_sum = (_fit_X[mask_ == 1, col]).sum().to(X.dtype)
+                #    div = torch.where(mask_sum > 0, mask_sum,
+                #                       torch.tensor([1], dtype=mask_sum.dtype))
+                #    X[all_nan_receivers_idx, col] = col_sum / div
+                #
+                #     # receivers with at least one defined distance
+                #     receivers_idx = receivers_idx[~all_nan_dist_mask]
+                #     dist_subset = dist_chunk
+                #           [dist_idx_map[receivers_idx]][:, potential_donors_idx]
+                # else
+                #     ... identity
+                X, dist_subset, receivers_idx = self._col_cond(
+                    X,
+                    dist_subset,
+                    mask_fit_X,
+                    _fit_X,
+                    receivers_idx,
+                    all_nan_receivers_idx,
+                    all_nan_dist_mask,
+                    dist_chunk,
+                    dist_idx_map,
+                    potential_donors_idx,
+                )
+
+                # when all_nan_receivers_idx is not empty (training set is big)
+                tn = torch.tensor(self.n_neighbors)
+                n_neighbors = torch.where(
+                    tn < potential_donors_idx.shape[0],
+                    tn,
+                    potential_donors_idx.shape[0],
+                )
+                # to make sure n_neighbors > 0
+                n_neighbors = torch.where(
+                    n_neighbors <= 0,
+                    torch.tensor([1], dtype=n_neighbors.dtype),
+                    n_neighbors,
+                )
+                value = self._calc_impute(
+                    dist_subset,
+                    n_neighbors,
+                    _fit_X[potential_donors_idx, col],
+                    mask_fit_X[potential_donors_idx, col],
+                )
+                X[receivers_idx, col] = value.to(X.dtype)
+                return X
+
+            def forward(
+                self,
+                X,
+                dist_chunk,
+                non_missing_fix_X,
+                mask_fit_X,
+                dist_idx_map,
+                mask,
+                row_missing_idx,
+                _fit_X,
+            ):
+                return self.process_one_col(
+                    X,
+                    dist_chunk,
+                    non_missing_fix_X,
+                    mask_fit_X,
+                    dist_idx_map,
+                    mask,
+                    row_missing_idx,
+                    _fit_X,
+                )
+
+        class MakeDictIdxMap(torch.nn.Module):
+            def forward(self, X, row_missing_idx):
+                dist_idx_map = torch.zeros(X.shape[0], dtype=int)
+                dist_idx_map[row_missing_idx] = torch.arange(row_missing_idx.shape[0])
+                return dist_idx_map
+
+        class TorchKNNImputer(torch.nn.Module):
+            def __init__(self, knn_imputer):
+                super().__init__()
+                assert (
+                    knn_imputer.metric == "nan_euclidean"
+                ), f"Not implemented for metric={knn_imputer.metric!r}"
+                self.dist = NanEuclidean()
+                cols = []
+                for col in range(knn_imputer._fit_X.shape[1]):
+                    cols.append(
+                        ColProcessor(col, knn_imputer.n_neighbors, knn_imputer.weights)
+                    )
+                self.columns = torch.nn.ModuleList(cols)
+                # refactoring
+                self._make_dict_idx_map = MakeDictIdxMap()
+                # knn imputer
+                self.missing_values = knn_imputer.missing_values
+                self.n_neighbors = knn_imputer.n_neighbors
+                self.weights = knn_imputer.weights
+                self.metric = knn_imputer.metric
+                self.keep_empty_features = knn_imputer.keep_empty_features
+                self.add_indicator = knn_imputer.add_indicator
+                # results of fitting
+                self.indicator_ = knn_imputer.indicator_
+                # The training results.
+                # self._fit_X = torch.from_numpy(knn_imputer._fit_X.astype(np.float32))
+                # self._mask_fit_X = torch.from_numpy(knn_imputer._mask_fit_X)
+                # self._valid_mask = torch.from_numpy(knn_imputer._valid_mask)
+
+            def _transform_indicator(self, X):
+                if self.add_indicator:
+                    if not hasattr(self, "indicator_"):
+                        raise ValueError(
+                            "Make sure to call _fit_indicator before _transform_indicator"
+                        )
+                    raise NotImplementedError(type(self.indicator_))
+                    # return self.indicator_.transform(X)
+                return None
+
+            def _concatenate_indicator(self, X_imputed, X_indicator):
+                if not self.add_indicator:
+                    return X_imputed
+                if X_indicator is None:
+                    raise ValueError(
+                        "Data from the missing indicator are not provided. Call "
+                        "_fit_indicator and _transform_indicator in the imputer "
+                        "implementation."
+                    )
+                return torch.cat([X_imputed, X_indicator], dim=0)
+
+            def transform(self, mask_fit_X, _valid_mask, _fit_X, X):
+                X = X.clone()
+                mask = _get_mask(X, self.missing_values)
+
+                X_indicator = self._transform_indicator(mask)
+
+                row_missing_idx = flatnonzero(mask[:, _valid_mask].any(axis=1))
+                non_missing_fix_X = torch.logical_not(mask_fit_X)
+
+                # Maps from indices from X to indices in dist matrix
+                dist_idx_map = self._make_dict_idx_map(X, row_missing_idx)
+
+                # process in fixed-memory chunks
+                pairwise_distances = self.dist(X[row_missing_idx, :], _fit_X)
+
+                # The export unfold the loop as it depends on the number of features.
+                # Fixed in this case.
+                for col_processor in self.columns:
+                    X = col_processor(
+                        X,
+                        pairwise_distances,
+                        non_missing_fix_X,
+                        mask_fit_X,
+                        dist_idx_map,
+                        mask,
+                        row_missing_idx,
+                        _fit_X,
+                    )
+
+                if self.keep_empty_features:
+                    Xc = X.clone()
+                    Xc[:, ~_valid_mask] = 0
+                else:
+                    Xc = X[:, _valid_mask]
+
+                return self._concatenate_indicator(Xc, X_indicator)
+
+            def forward(self, _mask_fit_X, _valid_mask, _fit_X, X):
+                return self.transform(_mask_fit_X, _valid_mask, _fit_X, X)
+
+        return TorchKNNImputer, NanEuclidean
+
     @unittest.skipIf(KNNImputer is None, reason="new in 0.22")
-    @unittest.skipIf(TARGET_OPSET < 9, reason="not available")
+    @unittest.skipIf(
+        pv.Version(ort_version) <= pv.Version("1.16.0"),
+        reason="onnxruntime not recent enough",
+    )
     @ignore_warnings(category=DeprecationWarning)
-    def test_sklearn_knn_imputer(self):
+    def test_sklearn_knn_imputer_main(self):
         x_train = numpy.array(
             [
                 [1, 2, numpy.nan, 12],
@@ -933,7 +1410,7 @@ class TestNearestNeighbourConverter(unittest.TestCase):
             dtype=numpy.float32,
         )
         model = KNNImputer(n_neighbors=3, metric="nan_euclidean").fit(x_train)
-        for opset in [TARGET_OPSET, 9, 10, 11, 12]:
+        for opset in [TARGET_OPSET, 18]:
             if opset > TARGET_OPSET:
                 continue
             model_onnx = convert_sklearn(
@@ -942,13 +1419,87 @@ class TestNearestNeighbourConverter(unittest.TestCase):
                 [("input", FloatTensorType((None, x_test.shape[1])))],
                 target_opset=opset,
             )
-            self.assertIsNotNone(model_onnx)
+
+            try:
+                import torch
+
+                has_torch = True
+            except ImportError:
+                has_torch = False
+
+            if has_torch:
+                # DEBUG
+                with open("debug.onnx", "wb") as f:
+                    f.write(model_onnx.SerializeToString())
+
+                from sklearn.metrics.pairwise import nan_euclidean_distances
+
+                tmodel_cls, dist_cls = self._get_torch_knn_imputer()
+                tmodel = tmodel_cls(model)
+                mm = dist_cls()
+                td = mm(
+                    torch.from_numpy(x_test),
+                    torch.from_numpy(model._fit_X.astype(numpy.float32)),
+                )
+                skl = nan_euclidean_distances(
+                    x_test, model._fit_X.astype(numpy.float32)
+                )
+                assert_almost_equal(td.numpy(), skl, decimal=3)
+
+                ty = tmodel.transform(
+                    torch.from_numpy(model._mask_fit_X),
+                    torch.from_numpy(model._valid_mask),
+                    torch.from_numpy(model._fit_X.astype(numpy.float32)),
+                    torch.from_numpy(x_test),
+                )
+                skl = model.transform(x_test)
+                assert_almost_equal(ty.numpy(), skl, decimal=3)
+
+                from experimental_experiment.reference import ExtendedReferenceEvaluator
+
+                ExtendedReferenceEvaluator(model_onnx, verbose=0).run(
+                    None, {"input": x_test}
+                )
+
+                # let's inline first
+                import onnx.inliner
+
+                inlined = onnx.inliner.inline_local_functions(model_onnx)
+                with open("debug_inlined.onnx", "wb") as f:
+                    f.write(inlined.SerializeToString())
+
+                import onnxruntime
+
+                opts = onnxruntime.SessionOptions()
+                # opts.log_severity_level = 0
+                # opts.log_verbosity_level = 0
+                # opts.graph_optimization_level = (
+                #     onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL
+                # )
+                sess = onnxruntime.InferenceSession(
+                    inlined.SerializeToString(),
+                    opts,
+                    providers=["CPUExecutionProvider"],
+                )
+                runopts = onnxruntime.RunOptions()
+                # runopts.log_severity_level = 0
+                # runopts.log_verbosity_level = 0
+                sess.run(None, {"input": x_test}, runopts)
+
             dump_data_and_model(
-                x_test, model, model_onnx, basename="SklearnKNNImputer%d" % opset
+                x_test,
+                model,
+                model_onnx,
+                basename="SklearnKNNImputer%d" % opset,
+                backend="onnxruntime",
+                verbose=0,
             )
 
     @unittest.skipIf(KNNImputer is None, reason="new in 0.22")
-    @unittest.skipIf(TARGET_OPSET < 9, reason="not available")
+    @unittest.skipIf(
+        pv.Version(ort_version) <= pv.Version("1.16.0"),
+        reason="onnxruntime not recent enough",
+    )
     @ignore_warnings(category=DeprecationWarning)
     def test_sklearn_knn_imputer_cdist(self):
         x_train = numpy.array(
@@ -975,7 +1526,7 @@ class TestNearestNeighbourConverter(unittest.TestCase):
                 options={id(model): {"optim2": "cdist"}},
             )
 
-        for opset in [TARGET_OPSET, 12, 11, 10, 9]:
+        for opset in [TARGET_OPSET, 18]:
             if opset > TARGET_OPSET:
                 continue
             model_onnx = convert_sklearn(
@@ -986,10 +1537,14 @@ class TestNearestNeighbourConverter(unittest.TestCase):
                 options={id(model): {"optim": "cdist"}},
             )
             self.assertIsNotNone(model_onnx)
-            self.assertIn('op_type: "cdist"', str(model_onnx).lower())
+            # self.assertIn('op_type: "cdist"', str(model_onnx).lower())
             self.assertNotIn("scan", str(model_onnx).lower())
             dump_data_and_model(
-                x_test, model, model_onnx, basename="SklearnKNNImputer%dcdist" % opset
+                x_test,
+                model,
+                model_onnx,
+                basename="SklearnKNNImputer%dcdist" % opset,
+                backend="onnxruntime",
             )
 
     @unittest.skipIf(
@@ -1123,6 +1678,58 @@ class TestNearestNeighbourConverter(unittest.TestCase):
             basename="SklearnKNeighborsClassifierMReg3-Out0",
         )
 
+    @unittest.skipIf(KNNImputer is None, reason="new in 0.22")
+    @unittest.skipIf(
+        pv.Version(ort_version) <= pv.Version("1.16.0"),
+        reason="onnxruntime not recent enough",
+    )
+    @ignore_warnings(category=DeprecationWarning)
+    @unittest.skipIf(
+        pv.Version(ort_version) < pv.Version("1.20.0"),
+        reason="onnxruntime not recent enough",
+    )
+    @unittest.skipIf(
+        sys.platform != "linux" and pv.Version(skl_version) < pv.Version("1.6.0"),
+        "investigate why topk returns different results",
+    )
+    def test_sklearn_knn_imputer_issue_2025(self):
+        # This test is about having nan or the fact TopK
+        # does not handle largest=1 in opset < 11.
+        from onnxruntime import InferenceSession
+
+        data = (numpy.arange(14) + 100).reshape((-1, 2)).astype(float)
+        data[:, 0] += 1000
+        for i in range(5):
+            data[i, i % 2] = numpy.nan
+        imputer = KNNImputer(n_neighbors=3, metric="nan_euclidean")
+        imputer.fit(data)
+        initial_type = [("float_input", FloatTensorType([None, data.shape[1]]))]
+        onnx_model = convert_sklearn(imputer, initial_types=initial_type)
+        input_data = data.astype(numpy.float32)
+        expected = imputer.transform(input_data)
+        got = ReferenceEvaluator(onnx_model, verbose=0).run(
+            None, {"float_input": input_data}
+        )[0]
+        assert_almost_equal(expected, got)
+
+        # in case onnruntime fails
+        # from experimental_experiment.reference import OrtEval
+
+        # got = OrtEval(onnx_model, verbose=10).run(None, {"float_input": input_data})[0]
+        # assert_almost_equal(expected, got)
+
+        got = InferenceSession(
+            onnx_model.SerializeToString(), providers=["CPUExecutionProvider"]
+        ).run(None, {"float_input": input_data})[0]
+        assert_almost_equal(expected, got)
+        dump_data_and_model(
+            input_data,
+            imputer,
+            onnx_model,
+            basename="SklearnKNNImputer2025",
+            backend="onnxruntime",
+        )
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)
diff --git a/tests/test_utils/utils_backend.py b/tests/test_utils/utils_backend.py
index ef4a13c..85cd7b4 100644
--- a/tests/test_utils/utils_backend.py
+++ b/tests/test_utils/utils_backend.py
@@ -324,6 +324,8 @@ def compare_outputs(expected, output, verbose=False, **kwargs):
                 else:
                     return OnnxRuntimeAssertionError(str(e))
         else:
+            if "Out0" in kwargs:
+                del kwargs["Out0"]
             try:
                 assert_array_almost_equal(expected, output, verbose=verbose, **kwargs)
             except Exception as e:
diff --git a/tests/test_utils/utils_backend_onnx.py b/tests/test_utils/utils_backend_onnx.py
index 914e985..627bd0b 100644
--- a/tests/test_utils/utils_backend_onnx.py
+++ b/tests/test_utils/utils_backend_onnx.py
@@ -323,7 +323,7 @@ if onnx_opset_version() >= 18:
         def __init__(self, *args, new_ops=None, **kwargs):
             # filter out new_ops
             onx = args[0]
-            if isinstance(onx, onnx.ModelProto):
+            if isinstance(onx, (onnx.ModelProto, onnx.FunctionProto, onnx.GraphProto)):
                 model = onx
             elif isinstance(onx, bytes):
                 model = onnx.load(io.BytesIO(onx))
@@ -333,9 +333,10 @@ if onnx_opset_version() >= 18:
             else:
                 raise TypeError(f"Not implemented for {type(args[0])}.")
             main_domain = None
-            for dom in model.opset_import:
-                if dom.domain == "":
-                    main_domain = dom.version
+            if hasattr(model, "opset_import"):
+                for dom in model.opset_import:
+                    if dom.domain == "":
+                        main_domain = dom.version
             if main_domain is None:
                 main_domain = 1
 
@@ -365,7 +366,8 @@ if onnx_opset_version() >= 18:
 
             self._main_domain = main_domain
             self._new_ops = new_new_ops
-            self._opset_import = model.opset_import
+            if hasattr(model, "opset_import"):
+                self._opset_import = model.opset_import
 
             # calls the constructor
             super().__init__(*args, new_ops=new_new_ops, **kwargs)
@@ -383,6 +385,7 @@ if onnx_opset_version() >= 18:
                 "opsets": self.opsets,
                 "verbose": self.verbose,
                 "new_ops": self.new_ops_,
+                "evaluator_cls": self.__class__,
             }
             if self.input_types_:
                 all_types = {i.name: i.type for i in self.onnx_graph_.input}