Fixes light API for operators with two outputs (#45)

* fix operators with two outputs

* version

* add more checks

Author: xadupre

CHANGELOGS.rst

@@ -1,6 +1,11 @@
 Change Logs
 ===========
 
+0.1.3
++++++
+
+* :pr:`45`: fixes light API for operators with two outputs
+
 0.1.2
 +++++
 

_unittests/ut_light_api/test_light_api.py

@@ -402,6 +402,45 @@ def test_operator_bool(self):
             got = ref.run(None, {"X": a, "Y": b})[0]
             self.assertEqualArray(f(a, b), got)
 
+    def test_topk(self):
+        onx = (
+            start()
+            .vin("X", np.float32)
+            .vin("K", np.int64)
+            .bring("X", "K")
+            .TopK()
+            .rename("Values", "Indices")
+            .vout()
+            .to_onnx()
+        )
+        self.assertIsInstance(onx, ModelProto)
+        ref = ReferenceEvaluator(onx)
+        x = np.array([[0, 1, 2, 3], [9, 8, 7, 6]], dtype=np.float32)
+        k = np.array([2], dtype=np.int64)
+        got = ref.run(None, {"X": x, "K": k})
+        self.assertEqualArray(np.array([[3, 2], [9, 8]], dtype=np.float32), got[0])
+        self.assertEqualArray(np.array([[3, 2], [0, 1]], dtype=np.int64), got[1])
+
+    def test_topk_reverse(self):
+        onx = (
+            start()
+            .vin("X", np.float32)
+            .vin("K", np.int64)
+            .bring("X", "K")
+            .TopK(largest=0)
+            .rename("Values", "Indices")
+            .vout()
+            .to_onnx()
+        )
+        self.assertIsInstance(onx, ModelProto)
+        ref = ReferenceEvaluator(onx)
+        x = np.array([[0, 1, 2, 3], [9, 8, 7, 6]], dtype=np.float32)
+        k = np.array([2], dtype=np.int64)
+        got = ref.run(None, {"X": x, "K": k})
+        self.assertEqualArray(np.array([[0, 1], [6, 7]], dtype=np.float32), got[0])
+        self.assertEqualArray(np.array([[0, 1], [3, 2]], dtype=np.int64), got[1])
+
 
 if __name__ == "__main__":
+    # TestLightApi().test_topk()
     unittest.main(verbosity=2)

onnx_array_api/__init__.py

@@ -3,5 +3,5 @@
 APIs to create ONNX Graphs.
 """
 
-__version__ = "0.1.2"
+__version__ = "0.1.3"
 __author__ = "Xavier Dupré"

onnx_array_api/light_api/_op_var.py

@@ -30,7 +30,7 @@ def ArgMin(
 
     def AveragePool(
         self,
-        auto_pad: str = b"NOTSET",
+        auto_pad: str = "NOTSET",
         ceil_mode: int = 0,
         count_include_pad: int = 0,
         dilations: Optional[List[int]] = None,
@@ -68,7 +68,7 @@ def Cast(self, saturate: int = 1, to: int = 0) -> "Var":
     def Celu(self, alpha: float = 1.0) -> "Var":
         return self.make_node("Celu", self, alpha=alpha)
 
-    def DepthToSpace(self, blocksize: int = 0, mode: str = b"DCR") -> "Var":
+    def DepthToSpace(self, blocksize: int = 0, mode: str = "DCR") -> "Var":
         return self.make_node("DepthToSpace", self, blocksize=blocksize, mode=mode)
 
     def DynamicQuantizeLinear(
@@ -137,7 +137,7 @@ def LpNormalization(self, axis: int = -1, p: int = 2) -> "Var":
 
     def LpPool(
         self,
-        auto_pad: str = b"NOTSET",
+        auto_pad: str = "NOTSET",
         ceil_mode: int = 0,
         dilations: Optional[List[int]] = None,
         kernel_shape: Optional[List[int]] = None,

onnx_array_api/light_api/_op_vars.py

@@ -6,7 +6,7 @@ class OpsVars:
     Operators taking multiple inputs.
     """
 
-    def BitShift(self, direction: str = b"") -> "Var":
+    def BitShift(self, direction: str = "") -> "Var":
         return self.make_node("BitShift", *self.vars_, direction=direction)
 
     def CenterCropPad(self, axes: Optional[List[int]] = None) -> "Var":
@@ -42,7 +42,7 @@ def Concat(self, axis: int = 0) -> "Var":
 
     def Conv(
         self,
-        auto_pad: str = b"NOTSET",
+        auto_pad: str = "NOTSET",
         dilations: Optional[List[int]] = None,
         group: int = 1,
         kernel_shape: Optional[List[int]] = None,
@@ -66,7 +66,7 @@ def Conv(
 
     def ConvInteger(
         self,
-        auto_pad: str = b"NOTSET",
+        auto_pad: str = "NOTSET",
         dilations: Optional[List[int]] = None,
         group: int = 1,
         kernel_shape: Optional[List[int]] = None,
@@ -90,7 +90,7 @@ def ConvInteger(
 
     def ConvTranspose(
         self,
-        auto_pad: str = b"NOTSET",
+        auto_pad: str = "NOTSET",
         dilations: Optional[List[int]] = None,
         group: int = 1,
         kernel_shape: Optional[List[int]] = None,
@@ -155,7 +155,7 @@ def DeformConv(
     def DequantizeLinear(self, axis: int = 1) -> "Var":
         return self.make_node("DequantizeLinear", *self.vars_, axis=axis)
 
-    def Einsum(self, equation: str = b"") -> "Var":
+    def Einsum(self, equation: str = "") -> "Var":
         return self.make_node("Einsum", *self.vars_, equation=equation)
 
     def Gather(self, axis: int = 0) -> "Var":
@@ -174,8 +174,8 @@ def Gemm(
     def GridSample(
         self,
         align_corners: int = 0,
-        mode: str = b"bilinear",
-        padding_mode: str = b"zeros",
+        mode: str = "bilinear",
+        padding_mode: str = "zeros",
     ) -> "Var":
         return self.make_node(
             "GridSample",
@@ -240,7 +240,7 @@ def Mod(self, fmod: int = 0) -> "Var":
         return self.make_node("Mod", *self.vars_, fmod=fmod)
 
     def NegativeLogLikelihoodLoss(
-        self, ignore_index: int = 0, reduction: str = b"mean"
+        self, ignore_index: int = 0, reduction: str = "mean"
     ) -> "Var":
         return self.make_node(
             "NegativeLogLikelihoodLoss",
@@ -257,12 +257,12 @@ def NonMaxSuppression(self, center_point_box: int = 0) -> "Var":
     def OneHot(self, axis: int = -1) -> "Var":
         return self.make_node("OneHot", *self.vars_, axis=axis)
 
-    def Pad(self, mode: str = b"constant") -> "Var":
+    def Pad(self, mode: str = "constant") -> "Var":
         return self.make_node("Pad", *self.vars_, mode=mode)
 
     def QLinearConv(
         self,
-        auto_pad: str = b"NOTSET",
+        auto_pad: str = "NOTSET",
         dilations: Optional[List[int]] = None,
         group: int = 1,
         kernel_shape: Optional[List[int]] = None,
@@ -431,13 +431,13 @@ def Resize(
         self,
         antialias: int = 0,
         axes: Optional[List[int]] = None,
-        coordinate_transformation_mode: str = b"half_pixel",
+        coordinate_transformation_mode: str = "half_pixel",
         cubic_coeff_a: float = -0.75,
         exclude_outside: int = 0,
         extrapolation_value: float = 0.0,
-        keep_aspect_ratio_policy: str = b"stretch",
-        mode: str = b"nearest",
-        nearest_mode: str = b"round_prefer_floor",
+        keep_aspect_ratio_policy: str = "stretch",
+        mode: str = "nearest",
+        nearest_mode: str = "round_prefer_floor",
     ) -> "Var":
         axes = axes or []
         return self.make_node(
@@ -456,8 +456,8 @@ def Resize(
 
     def RoiAlign(
         self,
-        coordinate_transformation_mode: str = b"half_pixel",
-        mode: str = b"avg",
+        coordinate_transformation_mode: str = "half_pixel",
+        mode: str = "avg",
         output_height: int = 1,
         output_width: int = 1,
         sampling_ratio: int = 0,
@@ -480,12 +480,12 @@ def STFT(self, onesided: int = 1) -> "Var":
     def Scatter(self, axis: int = 0) -> "Var":
         return self.make_node("Scatter", *self.vars_, axis=axis)
 
-    def ScatterElements(self, axis: int = 0, reduction: str = b"none") -> "Var":
+    def ScatterElements(self, axis: int = 0, reduction: str = "none") -> "Var":
         return self.make_node(
             "ScatterElements", *self.vars_, axis=axis, reduction=reduction
         )
 
-    def ScatterND(self, reduction: str = b"none") -> "Var":
+    def ScatterND(self, reduction: str = "none") -> "Var":
         return self.make_node("ScatterND", *self.vars_, reduction=reduction)
 
     def Slice(
@@ -498,13 +498,18 @@ def Slice(
 
     def TopK(self, axis: int = -1, largest: int = 1, sorted: int = 1) -> "Vars":
         return self.make_node(
-            "TopK", *self.vars_, axis=axis, largest=largest, sorted=sorted
+            "TopK",
+            *self.vars_,
+            axis=axis,
+            largest=largest,
+            sorted=sorted,
+            n_outputs=2,
         )
 
     def Trilu(self, upper: int = 1) -> "Var":
         return self.make_node("Trilu", *self.vars_, upper=upper)
 
-    def Upsample(self, mode: str = b"nearest") -> "Var":
+    def Upsample(self, mode: str = "nearest") -> "Var":
         return self.make_node("Upsample", *self.vars_, mode=mode)
 
     def Where(

onnx_array_api/light_api/model.py

@@ -28,8 +28,8 @@ class OnnxGraph:
     This API is meant to be light and allows the description of a graph.
 
     :param opset: main opset version
+    :param opsets: other opsets as a dictionary
     :param is_function: a :class:`onnx.ModelProto` or a :class:`onnx.FunctionProto`
-    :param opsets: others opsets as a dictionary
     """
 
     def __init__(