onnx 1.13.1__cp38-cp38-win32.whl → 1.14.1__cp38-cp38-win32.whl

onnx/reference/ops/aionnxml/op_tree_ensemble_helper.py

@@ -0,0 +1,109 @@
+# Copyright (c) ONNX Project Contributors
+
+# SPDX-License-Identifier: Apache-2.0
+# pylint: disable=R0911,R0913,R0914,W0221
+
+import numpy as np
+
+
+class TreeEnsembleAttributes:
+    def __init__(self):
+        self._names = []
+
+    def add(self, name, value):
+        if not name.endswith("_as_tensor"):
+            self._names.append(name)
+        if isinstance(value, list):
+            if name in {
+                "base_values",
+                "class_weights",
+                "nodes_values",
+                "nodes_hitrates",
+            }:
+                value = np.array(value, dtype=np.float32)
+            elif name.endswith("as_tensor"):
+                value = np.array(value)
+        setattr(self, name, value)
+
+    def __str__(self):
+        rows = ["Attributes"]
+        for name in self._names:
+            if name.endswith("_as_tensor"):
+                name = name.replace("_as_tensor", "")
+            rows.append(f"  {name}={getattr(self, name)}")
+        return "\n".join(rows)
+
+
+class TreeEnsemble:
+    def __init__(self, **kwargs):
+        self.atts = TreeEnsembleAttributes()
+
+        for name, value in kwargs.items():
+            self.atts.add(name, value)
+
+        self.tree_ids = list(sorted(set(self.atts.nodes_treeids)))  # type: ignore
+        self.root_index = {
+            tid: len(self.atts.nodes_treeids) for tid in self.tree_ids  # type: ignore
+        }
+        for index, tree_id in enumerate(self.atts.nodes_treeids):  # type: ignore
+            self.root_index[tree_id] = min(self.root_index[tree_id], index)
+        self.node_index = {
+            (tid, nid): i
+            for i, (tid, nid) in enumerate(
+                zip(self.atts.nodes_treeids, self.atts.nodes_nodeids)  # type: ignore
+            )
+        }
+
+    def __str__(self) -> str:
+        rows = ["TreeEnsemble", f"root_index={self.root_index}", str(self.atts)]
+        return "\n".join(rows)
+
+    def leaf_index_tree(self, X: np.ndarray, tree_id: int) -> int:
+        """
+        Computes the leaf index for one tree.
+        """
+        index = self.root_index[tree_id]
+        while self.atts.nodes_modes[index] != "LEAF":  # type: ignore
+            x = X[self.atts.nodes_featureids[index]]  # type: ignore
+            if np.isnan(x):
+                r = self.atts.nodes_missing_value_tracks_true[index] >= 1  # type: ignore
+            else:
+                rule = self.atts.nodes_modes[index]  # type: ignore
+                th = self.atts.nodes_values[index]  # type: ignore
+                if rule == "BRANCH_LEQ":
+                    r = x <= th
+                elif rule == "BRANCH_LT":
+                    r = x < th
+                elif rule == "BRANCH_GTE":
+                    r = x >= th
+                elif rule == "BRANCH_GT":
+                    r = x > th
+                elif rule == "BRANCH_EQ":
+                    r = x == th
+                elif rule == "BRANCH_NEQ":
+                    r = x != th
+                else:
+                    raise ValueError(
+                        f"Unexpected rule {rule!r} for node index {index}."
+                    )
+            nid = (
+                self.atts.nodes_truenodeids[index]  # type: ignore
+                if r
+                else self.atts.nodes_falsenodeids[index]  # type: ignore
+            )
+            index = self.node_index[tree_id, nid]
+        return index
+
+    def leave_index_tree(self, X: np.ndarray) -> np.ndarray:
+        """
+        Computes the leave index for all trees.
+        """
+        if len(X.shape) == 1:
+            X = X.reshape((1, -1))
+        outputs = []
+        for row in X:
+            outs = []
+            for tree_id in self.tree_ids:
+                outs.append(self.leaf_index_tree(row, tree_id))
+            outputs.append(outs)
+        return np.array(outputs)

onnx/reference/ops/aionnxml/op_tree_ensemble_regressor.py

@@ -0,0 +1,105 @@
+# Copyright (c) ONNX Project Contributors
+
+# SPDX-License-Identifier: Apache-2.0
+# pylint: disable=R0912,R0913,R0914,W0221
+
+import numpy as np
+
+from onnx.reference.ops.aionnxml._op_run_aionnxml import OpRunAiOnnxMl
+from onnx.reference.ops.aionnxml.op_tree_ensemble_helper import TreeEnsemble
+
+
+class TreeEnsembleRegressor(OpRunAiOnnxMl):
+    """
+    `nodes_hitrates` and `nodes_hitrates_as_tensor` are not used.
+    """
+
+    def _run(  # type: ignore
+        self,
+        X,
+        aggregate_function=None,
+        base_values=None,
+        base_values_as_tensor=None,
+        n_targets=None,
+        nodes_falsenodeids=None,
+        nodes_featureids=None,
+        nodes_hitrates=None,
+        nodes_hitrates_as_tensor=None,
+        nodes_missing_value_tracks_true=None,
+        nodes_modes=None,
+        nodes_nodeids=None,
+        nodes_treeids=None,
+        nodes_truenodeids=None,
+        nodes_values=None,
+        nodes_values_as_tensor=None,
+        post_transform=None,
+        target_ids=None,
+        target_nodeids=None,
+        target_treeids=None,
+        target_weights=None,
+        target_weights_as_tensor=None,
+    ):
+        nmv = nodes_missing_value_tracks_true
+        tr = TreeEnsemble(
+            base_values=base_values,
+            base_values_as_tensor=base_values_as_tensor,
+            nodes_falsenodeids=nodes_falsenodeids,
+            nodes_featureids=nodes_featureids,
+            nodes_hitrates=nodes_hitrates,
+            nodes_hitrates_as_tensor=nodes_hitrates_as_tensor,
+            nodes_missing_value_tracks_true=nmv,
+            nodes_modes=nodes_modes,
+            nodes_nodeids=nodes_nodeids,
+            nodes_treeids=nodes_treeids,
+            nodes_truenodeids=nodes_truenodeids,
+            nodes_values=nodes_values,
+            nodes_values_as_tensor=nodes_values_as_tensor,
+            target_weights=target_weights,
+            target_weights_as_tensor=target_weights_as_tensor,
+        )
+        # unused unless for debugging purposes
+        self._tree = tr  # pylint: disable=W0201
+        leaves_index = tr.leave_index_tree(X)
+        res = np.zeros((leaves_index.shape[0], n_targets), dtype=X.dtype)
+        n_trees = len(set(tr.atts.nodes_treeids))  # type: ignore
+
+        target_index = {}  # type: ignore
+        for i, (tid, nid) in enumerate(zip(target_treeids, target_nodeids)):
+            if (tid, nid) not in target_index:
+                target_index[tid, nid] = []
+            target_index[tid, nid].append(i)
+        for i in range(res.shape[0]):
+            indices = leaves_index[i]
+            t_index = [
+                target_index[nodes_treeids[i], nodes_nodeids[i]] for i in indices
+            ]
+            if aggregate_function in ("SUM", "AVERAGE"):
+                for its in t_index:
+                    for it in its:
+                        res[i, target_ids[it]] += tr.atts.target_weights[it]  # type: ignore
+            elif aggregate_function == "MIN":
+                res[i, :] = np.finfo(res.dtype).max
+                for its in t_index:
+                    for it in its:
+                        res[i, target_ids[it]] = min(
+                            res[i, target_ids[it]], tr.atts.target_weights[it]  # type: ignore
+                        )
+            elif aggregate_function == "MAX":
+                res[i, :] = np.finfo(res.dtype).min
+                for its in t_index:
+                    for it in its:
+                        res[i, target_ids[it]] = max(
+                            res[i, target_ids[it]], tr.atts.target_weights[it]  # type: ignore
+                        )
+            else:
+                raise NotImplementedError(
+                    f"aggregate_transform={aggregate_function!r} " f"not supported yet."
+                )
+        if aggregate_function == "AVERAGE":
+            res /= n_trees
+        if base_values is not None:
+            res[:, :] = np.array(base_values).reshape((1, -1))
+
+        if post_transform in (None, "NONE"):
+            return (res,)
+        raise NotImplementedError(f"post_transform={post_transform!r} not implemented.")

onnx/reference/ops/experimental/__init__.py

@@ -1,3 +1,5 @@
+# Copyright (c) ONNX Project Contributors
+#
 # SPDX-License-Identifier: Apache-2.0
 
-from ._op_list import load_op
+from onnx.reference.ops.experimental._op_list import load_op

onnx/reference/ops/experimental/_op_list.py

@@ -1,44 +1,21 @@
+# Copyright (c) ONNX Project Contributors
+
 # SPDX-License-Identifier: Apache-2.0
 # pylint: disable=C0415,R0912,W0611,W0603
 
 import textwrap
-from typing import Any, Union
-
-from onnx.reference.op_run import OpFunction, _split_class_name
+from typing import Any, Dict, Union
 
-from ._op_run_experimental import OpRunExperimental
-from .op_im2col import Im2Col
+from onnx.reference.op_run import OpFunction
+from onnx.reference.ops._helpers import build_registered_operators_any_domain
+from onnx.reference.ops.experimental._op_run_experimental import OpRunExperimental
+from onnx.reference.ops.experimental.op_im2col import Im2Col  # noqa: F401
 
 
-def _build_registered_operators():  # type: ignore
-    clo = globals().copy()
-    reg_ops = {}  # type: ignore
-    for class_name, class_type in clo.items():
-        if class_name[0] == "_" or class_name in {
-            "Any",
-            "cl",
-            "clo",
-            "class_name",
-            "textwrap",
-            "Union",
-        }:
-            continue  # pragma: no cover
-        if isinstance(class_type, type(load_op)):
-            continue
-        try:
-            issub = issubclass(class_type, OpRunExperimental)
-        except TypeError as e:
-            raise TypeError(
-                f"Unexpected variable type {class_type!r} and class_name={class_name!r}."
-            ) from e
-        if issub:
-            op_type, op_version = _split_class_name(class_name)
-            if op_type not in reg_ops:
-                reg_ops[op_type] = {}
-            reg_ops[op_type][op_version] = class_type
-    if len(reg_ops) == 0:
-        raise RuntimeError("No registered operators. The installation went wrong.")
-    return reg_ops
+def _build_registered_operators() -> (
+    Dict[str, Dict[Union[int, None], OpRunExperimental]]
+):
+    return build_registered_operators_any_domain(globals().copy())  # type: ignore[return-value]
 
 
 def load_op(
@@ -55,7 +32,7 @@ def load_op(
     """
     global _registered_operators
     if _registered_operators is None:
-        _registered_operators = _build_registered_operators()
+        _registered_operators = _build_registered_operators()  # type: ignore[assignment]
     if custom is not None:
         return lambda *args: OpFunction(*args, impl=custom)  # type: ignore
     if domain != "experimental":
@@ -98,4 +75,6 @@ def load_op(
     return cl
 
 
-_registered_operators = None
+# Python 3.7 does not support this annotation for a global variable.
+# _registered_operators: TOptional[Dict[str, Dict[Union[int, None], OpRunExperimental]]] = None
+_registered_operators = None  # type: ignore

onnx/reference/ops/experimental/_op_run_experimental.py

@@ -1,3 +1,5 @@
+# Copyright (c) ONNX Project Contributors
+#
 # SPDX-License-Identifier: Apache-2.0
 # pylint: disable=R0913,W0221
 

onnx/reference/ops/experimental/op_im2col.py

@@ -1,10 +1,11 @@
+# Copyright (c) ONNX Project Contributors
+#
 # SPDX-License-Identifier: Apache-2.0
 # pylint: disable=R0913,R0914,W0221
 
+from onnx.reference.ops.experimental._op_run_experimental import OpRunExperimental
 from onnx.reference.ops.op_conv import im2col
 
-from ._op_run_experimental import OpRunExperimental
-
 
 class Im2Col(OpRunExperimental):
     def _run(self, img, kernel_shape, dilations=None, pads=None, strides=None):  # type: ignore

onnx/reference/ops/op_abs.py

@@ -1,9 +1,11 @@
+# Copyright (c) ONNX Project Contributors
+
 # SPDX-License-Identifier: Apache-2.0
 # pylint: disable=W0221
 
 import numpy as np
 
-from ._op import OpRunUnaryNum
+from onnx.reference.ops._op import OpRunUnaryNum
 
 
 class Abs(OpRunUnaryNum):

onnx/reference/ops/op_acos.py

@@ -1,9 +1,11 @@
+# Copyright (c) ONNX Project Contributors
+
 # SPDX-License-Identifier: Apache-2.0
 # pylint: disable=W0221
 
 import numpy as np
 
-from ._op import OpRunUnaryNum
+from onnx.reference.ops._op import OpRunUnaryNum
 
 
 class Acos(OpRunUnaryNum):

onnx/reference/ops/op_acosh.py

@@ -1,9 +1,11 @@
+# Copyright (c) ONNX Project Contributors
+
 # SPDX-License-Identifier: Apache-2.0
 # pylint: disable=W0221
 
 import numpy as np
 
-from ._op import OpRunUnaryNum
+from onnx.reference.ops._op import OpRunUnaryNum
 
 
 class Acosh(OpRunUnaryNum):

onnx/reference/ops/op_add.py

@@ -1,9 +1,11 @@
+# Copyright (c) ONNX Project Contributors
+
 # SPDX-License-Identifier: Apache-2.0
 # pylint: disable=W0221
 
 import numpy as np
 
-from ._op import OpRunBinaryNumpy
+from onnx.reference.ops._op import OpRunBinaryNumpy
 
 
 class Add(OpRunBinaryNumpy):

onnx/reference/ops/op_and.py

@@ -1,9 +1,11 @@
+# Copyright (c) ONNX Project Contributors
+
 # SPDX-License-Identifier: Apache-2.0
 # pylint: disable=W0221
 
 import numpy as np
 
-from ._op import OpRunBinary
+from onnx.reference.ops._op import OpRunBinary
 
 
 class And(OpRunBinary):

onnx/reference/ops/op_argmax.py

@@ -1,9 +1,10 @@
+# Copyright (c) ONNX Project Contributors
+
 # SPDX-License-Identifier: Apache-2.0
 # pylint: disable=W0221
 
 import numpy as np
 
-from onnx.defs import onnx_opset_version
 from onnx.reference.op_run import OpRun
 
 
@@ -28,20 +29,14 @@ class _ArgMax(OpRun):
         return (_argmax(data, axis=axis, keepdims=keepdims),)
 
 
-class ArgMax_11(_ArgMax):
+class ArgMax_1(_ArgMax):
     pass
 
 
 class ArgMax_12(_ArgMax):
     def _run(self, data, axis=None, keepdims=None, select_last_index=None):  # type: ignore
         if select_last_index == 0:  # type: ignore
-            return _ArgMax._run(self, data)
+            return _ArgMax._run(self, data, axis=axis, keepdims=keepdims)
         return (
             _argmax_use_numpy_select_last_index(data, axis=axis, keepdims=keepdims),
         )
-
-
-if onnx_opset_version() >= 12:
-    ArgMax = ArgMax_12
-else:
-    ArgMax = ArgMax_11  # type: ignore

onnx/reference/ops/op_argmin.py

@@ -1,9 +1,10 @@
+# Copyright (c) ONNX Project Contributors
+
 # SPDX-License-Identifier: Apache-2.0
 # pylint: disable=W0221
 
 import numpy as np
 
-from onnx.defs import onnx_opset_version
 from onnx.reference.op_run import OpRun
 
 
@@ -28,20 +29,14 @@ class _ArgMin(OpRun):
         return (_argmin(data, axis=axis, keepdims=keepdims),)
 
 
-class ArgMin_11(_ArgMin):
+class ArgMin_1(_ArgMin):
     pass
 
 
 class ArgMin_12(_ArgMin):
     def _run(self, data, axis=None, keepdims=None, select_last_index=None):  # type: ignore
         if select_last_index == 0:  # type: ignore
-            return _ArgMin._run(self, data)
+            return _ArgMin._run(self, data, axis=axis, keepdims=keepdims)
         return (
             _argmin_use_numpy_select_last_index(data, axis=axis, keepdims=keepdims),
         )
-
-
-if onnx_opset_version() >= 12:
-    ArgMin = ArgMin_12
-else:
-    ArgMin = ArgMin_11  # type: ignore

onnx/reference/ops/op_asin.py

@@ -1,9 +1,11 @@
+# Copyright (c) ONNX Project Contributors
+
 # SPDX-License-Identifier: Apache-2.0
 # pylint: disable=W0221
 
 import numpy as np
 
-from ._op import OpRunUnaryNum
+from onnx.reference.ops._op import OpRunUnaryNum
 
 
 class Asin(OpRunUnaryNum):

onnx/reference/ops/op_asinh.py

@@ -1,9 +1,11 @@
+# Copyright (c) ONNX Project Contributors
+
 # SPDX-License-Identifier: Apache-2.0
 # pylint: disable=W0221
 
 import numpy as np
 
-from ._op import OpRunUnaryNum
+from onnx.reference.ops._op import OpRunUnaryNum
 
 
 class Asinh(OpRunUnaryNum):

onnx/reference/ops/op_atan.py

@@ -1,9 +1,11 @@
+# Copyright (c) ONNX Project Contributors
+
 # SPDX-License-Identifier: Apache-2.0
 # pylint: disable=W0221
 
 import numpy as np
 
-from ._op import OpRunUnaryNum
+from onnx.reference.ops._op import OpRunUnaryNum
 
 
 class Atan(OpRunUnaryNum):

onnx/reference/ops/op_atanh.py

@@ -1,9 +1,11 @@
+# Copyright (c) ONNX Project Contributors
+
 # SPDX-License-Identifier: Apache-2.0
 # pylint: disable=W0221
 
 import numpy as np
 
-from ._op import OpRunUnaryNum
+from onnx.reference.ops._op import OpRunUnaryNum
 
 
 class Atanh(OpRunUnaryNum):

onnx/reference/ops/op_attribute_has_value.py

@@ -1,3 +1,5 @@
+# Copyright (c) ONNX Project Contributors
+
 # SPDX-License-Identifier: Apache-2.0
 # pylint: disable=R0913,R0914,W0221,W0613
 

onnx/reference/ops/op_average_pool.py

@@ -1,10 +1,62 @@
+# Copyright (c) ONNX Project Contributors
+
 # SPDX-License-Identifier: Apache-2.0
 # pylint: disable=W0221,R0913,R0914
 
-from ._op_common_pool import CommonPool
+from onnx.reference.ops._op_common_pool import CommonPool
+
+
+class AveragePool_1(CommonPool):
+    def _run(  # type: ignore
+        self,
+        x,
+        auto_pad=None,
+        ceil_mode=None,
+        kernel_shape=None,
+        pads=None,
+        strides=None,
+        count_include_pad=None,
+    ):
+        return CommonPool._run(
+            self,
+            "AVG",
+            count_include_pad,
+            x,
+            auto_pad=auto_pad,
+            ceil_mode=ceil_mode,
+            dilations=None,
+            kernel_shape=kernel_shape,
+            pads=pads,
+            strides=strides,
+        )
+
+
+class AveragePool_7(CommonPool):
+    def _run(  # type: ignore
+        self,
+        x,
+        auto_pad=None,
+        ceil_mode=None,
+        kernel_shape=None,
+        pads=None,
+        strides=None,
+        count_include_pad=None,
+    ):
+        return CommonPool._run(
+            self,
+            "AVG",
+            count_include_pad,
+            x,
+            auto_pad=auto_pad,
+            ceil_mode=ceil_mode,
+            dilations=None,
+            kernel_shape=kernel_shape,
+            pads=pads,
+            strides=strides,
+        )
 
 
-class AveragePool(CommonPool):
+class AveragePool_11(CommonPool):
     def _run(  # type: ignore
         self,
         x,
@@ -27,3 +79,29 @@ class AveragePool(CommonPool):
             pads=pads,
             strides=strides,
         )
+
+
+class AveragePool_19(CommonPool):
+    def _run(  # type: ignore
+        self,
+        x,
+        auto_pad=None,
+        ceil_mode=None,
+        dilations=None,
+        kernel_shape=None,
+        pads=None,
+        strides=None,
+        count_include_pad=None,
+    ):
+        return CommonPool._run(
+            self,
+            "AVG",
+            count_include_pad,
+            x,
+            auto_pad=auto_pad,
+            ceil_mode=ceil_mode,
+            dilations=dilations,
+            kernel_shape=kernel_shape,
+            pads=pads,
+            strides=strides,
+        )

onnx/reference/ops/op_batch_normalization.py

@@ -1,9 +1,10 @@
+# Copyright (c) ONNX Project Contributors
+
 # SPDX-License-Identifier: Apache-2.0
 # pylint: disable=W0221,R0913,W0613
 
 import numpy as np
 
-from onnx.defs import onnx_opset_version
 from onnx.reference.op_run import OpRun
 
 
@@ -61,8 +62,18 @@ class BatchNormalization_6(OpRun):
 
 
 class BatchNormalization_9(OpRun):
-    def _run(self, x, scale, bias, mean, var, epsilon=None):  # type: ignore
-        res = _batchnorm_test_mode(x, scale, bias, mean, var, epsilon=epsilon)
+    def _run(self, x, scale, bias, mean, var, epsilon=None, momentum=None):  # type: ignore
+        if momentum is None:
+            res = _batchnorm_test_mode(x, scale, bias, mean, var, epsilon=epsilon)
+            return (res,)
+        axis = tuple(np.delete(np.arange(len(x.shape)), 1))
+        saved_mean = x.mean(axis=axis)
+        saved_var = x.var(axis=axis)
+        output_mean = mean * momentum + saved_mean * (1 - momentum)
+        output_var = var * momentum + saved_var * (1 - momentum)
+        res = _batchnorm_test_mode(
+            x, scale, bias, output_mean, output_var, epsilon=epsilon
+        )
         return (res,)
 
 
@@ -77,11 +88,3 @@ class BatchNormalization_14(OpRun):
             x, scale, bias, mean, var, momentum, epsilon
         )
         return res, output_mean, output_var
-
-
-if onnx_opset_version() >= 14:
-    BatchNormalization = BatchNormalization_14
-elif onnx_opset_version() >= 9:
-    BatchNormalization = BatchNormalization_9  # type: ignore
-else:
-    BatchNormalization = BatchNormalization_6  # type: ignore

onnx/reference/ops/op_bernoulli.py

@@ -1,9 +1,10 @@
+# Copyright (c) ONNX Project Contributors
+
 # SPDX-License-Identifier: Apache-2.0
 # pylint: disable=W0221
 
 from onnx.helper import np_dtype_to_tensor_dtype
-
-from ._op_common_random import _CommonRandom
+from onnx.reference.ops._op_common_random import _CommonRandom
 
 
 class Bernoulli(_CommonRandom):

onnx/reference/ops/op_bitshift.py

@@ -1,9 +1,11 @@
+# Copyright (c) ONNX Project Contributors
+
 # SPDX-License-Identifier: Apache-2.0
 # pylint: disable=W0221
 
 import numpy as np
 
-from ._op import OpRunBinaryNumpy
+from onnx.reference.ops._op import OpRunBinaryNumpy
 
 
 class BitShift(OpRunBinaryNumpy):