From 3a15ea15409d995e352d288a1eeccfbfb2751b8d Mon Sep 17 00:00:00 2001
From: Yuantao Feng <yuantao.feng@outlook.com>
Date: Tue, 13 Aug 2024 16:44:32 +0800
Subject: [PATCH] add models and data converted from onnx conformance

---
 testdata/dnn/onnx/data/input_top_k.pb         | Bin 0 -> 56 bytes
 .../onnx/data/input_top_k_negative_axis.pb    | Bin 0 -> 56 bytes
 .../dnn/onnx/data/input_top_k_smallest.pb     | Bin 0 -> 56 bytes
 testdata/dnn/onnx/data/output_top_k_0.pb      | Bin 0 -> 44 bytes
 testdata/dnn/onnx/data/output_top_k_1.pb      | Bin 0 -> 80 bytes
 .../onnx/data/output_top_k_negative_axis_0.pb | Bin 0 -> 44 bytes
 .../onnx/data/output_top_k_negative_axis_1.pb | Bin 0 -> 80 bytes
 .../dnn/onnx/data/output_top_k_smallest_0.pb  | Bin 0 -> 44 bytes
 .../dnn/onnx/data/output_top_k_smallest_1.pb  | Bin 0 -> 80 bytes
 .../generate_onnx_models_with_onnxscript.py   |  71 +++++++++++++++---
 testdata/dnn/onnx/models/top_k.onnx           | Bin 0 -> 214 bytes
 .../dnn/onnx/models/top_k_negative_axis.onnx  | Bin 0 -> 237 bytes
 testdata/dnn/onnx/models/top_k_smallest.onnx  | Bin 0 -> 254 bytes
 13 files changed, 62 insertions(+), 9 deletions(-)
 create mode 100644 testdata/dnn/onnx/data/input_top_k.pb
 create mode 100644 testdata/dnn/onnx/data/input_top_k_negative_axis.pb
 create mode 100644 testdata/dnn/onnx/data/input_top_k_smallest.pb
 create mode 100644 testdata/dnn/onnx/data/output_top_k_0.pb
 create mode 100644 testdata/dnn/onnx/data/output_top_k_1.pb
 create mode 100644 testdata/dnn/onnx/data/output_top_k_negative_axis_0.pb
 create mode 100644 testdata/dnn/onnx/data/output_top_k_negative_axis_1.pb
 create mode 100644 testdata/dnn/onnx/data/output_top_k_smallest_0.pb
 create mode 100644 testdata/dnn/onnx/data/output_top_k_smallest_1.pb
 create mode 100644 testdata/dnn/onnx/models/top_k.onnx
 create mode 100644 testdata/dnn/onnx/models/top_k_negative_axis.onnx
 create mode 100644 testdata/dnn/onnx/models/top_k_smallest.onnx

diff --git a/testdata/dnn/onnx/data/input_top_k.pb b/testdata/dnn/onnx/data/input_top_k.pb
new file mode 100644
index 0000000000000000000000000000000000000000..9b427a248fa32bc060984e5d313d8ccff78d6747
GIT binary patch
literal 56
ucmd;J=3o(E^fF)og9dvb<p4xL+yKN2fcO9qKLBC|M+OD~AXWfk14jV3s|Q8^

literal 0
HcmV?d00001

diff --git a/testdata/dnn/onnx/data/input_top_k_negative_axis.pb b/testdata/dnn/onnx/data/input_top_k_negative_axis.pb
new file mode 100644
index 0000000000000000000000000000000000000000..9b427a248fa32bc060984e5d313d8ccff78d6747
GIT binary patch
literal 56
ucmd;J=3o(E^fF)og9dvb<p4xL+yKN2fcO9qKLBC|M+OD~AXWfk14jV3s|Q8^

literal 0
HcmV?d00001

diff --git a/testdata/dnn/onnx/data/input_top_k_smallest.pb b/testdata/dnn/onnx/data/input_top_k_smallest.pb
new file mode 100644
index 0000000000000000000000000000000000000000..c707240c0985a513c428b48e56c1a6a10f1a1576
GIT binary patch
literal 56
ucmd;J=3o(E^fF)og9dvb<p4xL+yKN2fcO9qKLBC_M+OE3AQk{(21fw9D+fjZ

literal 0
HcmV?d00001

diff --git a/testdata/dnn/onnx/data/output_top_k_0.pb b/testdata/dnn/onnx/data/output_top_k_0.pb
new file mode 100644
index 0000000000000000000000000000000000000000..b021689446d811ce7a306bfe7a70759fe6342ebd
GIT binary patch
literal 44
scmd;J=3o|J^ipA9aByG%5)2Ft_6!UUfb;<%UI4@fjtmS6KrG-00Dh$gH~;_u

literal 0
HcmV?d00001

diff --git a/testdata/dnn/onnx/data/output_top_k_1.pb b/testdata/dnn/onnx/data/output_top_k_1.pb
new file mode 100644
index 0000000000000000000000000000000000000000..4158cd01cf9fb6e1bdc2c75b115da18f849c3ba4
GIT binary patch
literal 80
acmd;J=3o|J_wrz7fB+^a%?PDQ6$b!1%>dE>

literal 0
HcmV?d00001

diff --git a/testdata/dnn/onnx/data/output_top_k_negative_axis_0.pb b/testdata/dnn/onnx/data/output_top_k_negative_axis_0.pb
new file mode 100644
index 0000000000000000000000000000000000000000..b021689446d811ce7a306bfe7a70759fe6342ebd
GIT binary patch
literal 44
scmd;J=3o|J^ipA9aByG%5)2Ft_6!UUfb;<%UI4@fjtmS6KrG-00Dh$gH~;_u

literal 0
HcmV?d00001

diff --git a/testdata/dnn/onnx/data/output_top_k_negative_axis_1.pb b/testdata/dnn/onnx/data/output_top_k_negative_axis_1.pb
new file mode 100644
index 0000000000000000000000000000000000000000..4158cd01cf9fb6e1bdc2c75b115da18f849c3ba4
GIT binary patch
literal 80
acmd;J=3o|J_wrz7fB+^a%?PDQ6$b!1%>dE>

literal 0
HcmV?d00001

diff --git a/testdata/dnn/onnx/data/output_top_k_smallest_0.pb b/testdata/dnn/onnx/data/output_top_k_smallest_0.pb
new file mode 100644
index 0000000000000000000000000000000000000000..7a3d13db91ffa0d017b5bf04db0427efb3e24ca1
GIT binary patch
literal 44
ocmd;J=3o|J^ip8}g9dvb<-owu0K^M`_y7<yI5IE@0I`B20A&FMC;$Ke

literal 0
HcmV?d00001

diff --git a/testdata/dnn/onnx/data/output_top_k_smallest_1.pb b/testdata/dnn/onnx/data/output_top_k_smallest_1.pb
new file mode 100644
index 0000000000000000000000000000000000000000..6110e6409befbee57263f083fc7b95eca52457d1
GIT binary patch
literal 80
ecmd;J=3o|J_wrzX0!Ap!1QWm^&WxrGrVju)RRGKY

literal 0
HcmV?d00001

diff --git a/testdata/dnn/onnx/generate_onnx_models_with_onnxscript.py b/testdata/dnn/onnx/generate_onnx_models_with_onnxscript.py
index 7240e044c..bc85b8fde 100644
--- a/testdata/dnn/onnx/generate_onnx_models_with_onnxscript.py
+++ b/testdata/dnn/onnx/generate_onnx_models_with_onnxscript.py
@@ -19,7 +19,7 @@ def make_model_and_data(model, *args, **kwargs):
     name = model._name
 
     # TODO: support multiple outputs
-    output = model(*args) # eager mode
+    outputs = model(*args) # eager mode
 
     # Save model
     model_proto = model.to_model_proto()
@@ -34,7 +34,7 @@ def make_model_and_data(model, *args, **kwargs):
         model_proto_ = onnx.shape_inference.infer_shapes(model_proto)
         onnx.save(model_proto_, save_path)
 
-    # Save inputs and output
+    # Save inputs
     inputs = args
     if "force_saving_input_as_dtype_float32" in kwargs and kwargs["force_saving_input_as_dtype_float32"]:
         inputs = []
@@ -42,15 +42,41 @@ def make_model_and_data(model, *args, **kwargs):
             inputs.append(input.astype(np.float32))
     if len(args) == 1:
         input_file = os.path.join("data", "input_" + name)
-        np.save(input_file, inputs[0])
+        if "save_inputs_as_pb" in kwargs and kwargs["save_inputs_as_pb"]:
+            input_tensor = onnx.numpy_helper.from_array(inputs[0])
+            onnx.save_tensor(input_tensor, input_file + ".pb")
+        else:
+            np.save(input_file, inputs[0])
     else:
         for idx, input in enumerate(inputs, start=0):
-            input_files = os.path.join("data", "input_" + name + "_" + str(idx))
-            np.save(input_files, input)
-    if "force_saving_output_as_dtype_float32" in kwargs and kwargs["force_saving_output_as_dtype_float32"]:
-        output = output.astype(np.float32)
-    output_files = os.path.join("data", "output_" + name)
-    np.save(output_files, output)
+            input_file = os.path.join("data", "input_{}_{}".format(name, idx))
+            if "save_inputs_as_pb" in kwargs and kwargs["save_inputs_as_pb"]:
+                input_tensor = onnx.numpy_helper.from_array(input)
+                onnx.save_tensor(input_tensor, input_file + ".pb")
+            else:
+                np.save(input_file, input)
+
+    # Save outputs
+    if isinstance(outputs, tuple):
+        for idx, output in enumerate(outputs):
+            output_filepath = os.path.join("data", "output_{}_{}".format(name, idx))
+            if "force_saving_output_as_dtype_float32" in kwargs and kwargs["force_saving_output_as_dtype_float32"]:
+                output = output.astype(np.float32)
+            if "save_outputs_as_pb" in kwargs and kwargs["save_outputs_as_pb"]:
+                input_tensor = onnx.numpy_helper.from_array(output)
+                onnx.save_tensor(input_tensor, output_filepath + ".pb")
+            else:
+                np.save(output_filepath, output)
+    else:
+        output = outputs
+        if "force_saving_output_as_dtype_float32" in kwargs and kwargs["force_saving_output_as_dtype_float32"]:
+            output = output.astype(np.float32)
+        output_filepath = os.path.join("data", "output_" + name)
+        if "save_outputs_as_pb" in kwargs and kwargs["save_outputs_as_pb"]:
+            output_tensor = onnx.numpy_helper.from_array(output)
+            onnx.save_tensor(output_tensor, output_filepath + ".pb")
+        else:
+            np.save(output_filepath, output)
 
 '''
     It builds a model with two Gather ops sharing a single same indices:
@@ -390,3 +416,30 @@ def clip_div_shared_constant(x: ost.FLOAT[1, 8, 12, 10]) -> ost.FLOAT[1, 8, 12,
 def matmul_bcast(x: ost.FLOAT[64, 1, 16]) -> ost.FLOAT[64, 1, 8]:
     return op.MatMul(x, op.Constant(value=onnx.numpy_helper.from_array(B)))
 make_model_and_data(matmul_bcast, np.random.randn(64, 1, 16).astype(np.float32))
+
+''' TopK conformance
+'''
+
+top_k_K_arr = np.array([3], dtype=np.int64)
+@ost.script()
+def top_k(x: ost.FLOAT[3, 4]) -> (ost.FLOAT[3, 3], ost.INT64[3, 3]):
+    values, indices = op.TopK(x, op.Constant(value=onnx.numpy_helper.from_array(top_k_K_arr)), axis=1)
+    return values, indices
+
+@ost.script()
+def top_k_negative_axis(x: ost.FLOAT[3, 4]) -> (ost.FLOAT[3, 3], ost.INT64[3, 3]):
+    values, indices = op.TopK(x, op.Constant(value=onnx.numpy_helper.from_array(top_k_K_arr)), axis=-1)
+    return values, indices
+
+@ost.script()
+def top_k_smallest(x: ost.FLOAT[3, 4]) -> (ost.FLOAT[3, 3], ost.INT64[3, 3]):
+    values, indices = op.TopK(x, op.Constant(value=onnx.numpy_helper.from_array(top_k_K_arr)), axis=1, largest=0, sorted=1)
+    return values, indices
+
+top_k_input0 = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
+top_k_input1 = [0, 1, 2, 3, 4, 5, 6, 7, 11, 10, 9, 8]
+top_k_input0_arr = np.array(top_k_input0, dtype=np.float32).reshape(3, 4)
+top_k_input1_arr = np.array(top_k_input1, dtype=np.float32).reshape(3, 4)
+make_model_and_data(top_k,               top_k_input0_arr, save_inputs_as_pb=True, save_outputs_as_pb=True)
+make_model_and_data(top_k_negative_axis, top_k_input0_arr, save_inputs_as_pb=True, save_outputs_as_pb=True)
+make_model_and_data(top_k_smallest,      top_k_input1_arr, save_inputs_as_pb=True, save_outputs_as_pb=True)
diff --git a/testdata/dnn/onnx/models/top_k.onnx b/testdata/dnn/onnx/models/top_k.onnx
new file mode 100644
index 0000000000000000000000000000000000000000..9d00cd4ec913d89408d9396be01f47201a3f0c08
GIT binary patch
literal 214
zcmY+8y$*sf6op$_?Dc90mIWCZnZ%J%H(i}v3@8Ks6p~iM<M2Mdlof)BH#yTe-<J#1
z>`gJ1%siUNQL7lXwbCY6CW+8}<@aZqctCxz14f4HD9;G4k+%4y<fLC|DJpfT&ZV}R
zH0Fo8*(WpPd0Xj_cGSs;bWPoyZpXl)(!)7W3FrW`;YEm}=f7da5=V0ni66wSi{S%A
N03o^c&<#i82Ol+9CFTGC

literal 0
HcmV?d00001

diff --git a/testdata/dnn/onnx/models/top_k_negative_axis.onnx b/testdata/dnn/onnx/models/top_k_negative_axis.onnx
new file mode 100644
index 0000000000000000000000000000000000000000..403a32ab2f15c4e9069e1346f297152b83c08f0d
GIT binary patch
literal 237
zcmY+8yAFat5Jg#bm2qMSE*3~hNhK|I+G%TLLqInM9}5X9;dA&le?%7uCf;PKx#!Hl
zEPfM=I91QmH|08jT_G!#N)`Gzsnf?xgf<Yz+5<I)w1mY5ZjdzCQGD9BDsJWSlAlGT
zHBrEhMR^F9$kHaSyw4cY5;HcqtBNwYCsJHfmDeKaN+(wvxs7um9MA-+hh-j)`u`2h
aRyZE)A?pXx+hTYO5y0Y1d&rid;kaLx$u!yk

literal 0
HcmV?d00001

diff --git a/testdata/dnn/onnx/models/top_k_smallest.onnx b/testdata/dnn/onnx/models/top_k_smallest.onnx
new file mode 100644
index 0000000000000000000000000000000000000000..12d33c2cce4b81339234806b1edfa01e2fc4e950
GIT binary patch
literal 254
zcmY*Ty$*sf5N>Iu=cxqCf{u<%x-#mflZlgysX`2jKx{QU4)5p-SRtBt$v0fSU+Bj#
ziVMlit(rQ;d<Lts&_)#|nqu#%@<$qZKz*?WM!IV#j|py(*7$Krw^eB=vLeY&skSkT
z8Q+!FCK@1Db*4kwP$%+G<m!G-wF$}RAEB$uyGfIlj*0Y4SsktieN%aERr|nJNDqfV
nC7=V$f*(Wt_Wd*LIKl6N)+Bx=b{Tg|5dnncHis@b5<hqYwD~cF

literal 0
HcmV?d00001