support topk-10

Author: zhijxu-MS

tests/test_backend.py

@@ -1333,13 +1333,23 @@ def test_cancel_transpose(self):
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
     @check_opset_min_version(6, "cast")
-    def test_topk(self):
+    def test_topk1(self):
         x_val = np.arange(3 * 2 * 3).astype("float32")
         x = tf.placeholder(tf.float32, x_val.shape, name=_TFINPUT)
         values, _ = tf.nn.top_k(x, 5, sorted=True)
         _ = tf.identity(values, name=_TFOUTPUT)
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
+    @check_opset_min_version(10, "TopK with dynamic K")
+    def test_topk2(self):
+        x_val = np.arange(3 * 2 * 3).astype("float32")
+        x = tf.placeholder(tf.float32, x_val.shape, name=_TFINPUT)
+        k_val = np.array(10).astype(np.int32)
+        k = tf.placeholder(tf.int32, name=_TFINPUT1)
+        values, _ = tf.nn.top_k(x, k, sorted=True)
+        _ = tf.identity(values, name=_TFOUTPUT)
+        self._run_test_case([_OUTPUT], {_INPUT: x_val, _INPUT1: k_val})
+
     def test_stack_axis(self):
         for axis in [0, 1]:
             tf.reset_default_graph()

tf2onnx/onnx_opset/tensor.py

@@ -609,7 +609,7 @@ def version_7(cls, ctx, node, **kwargs):
         pass
 
 
-@tf_op("TopKV2")
+@tf_op("TopKV2", onnx_op="TopK")
 class TopKV2:
     @classmethod
     def version_4(cls, ctx, node, **kwargs):
@@ -634,6 +634,15 @@ def version_4(cls, ctx, node, **kwargs):
                                       name=new_cast_name, attr={"to": onnx_pb.TensorProto.INT32},
                                       shapes=[shapes[1]], dtypes=[onnx_pb.TensorProto.INT32])
 
+    @classmethod
+    def version_10(cls, ctx, node, **kwargs):
+        # onnx only supports input K as a 1D tesor with dtype int64
+        # while in tf, K is a 0D tensor with dtype int32
+        k_0d = node.input[1]
+        cast = ctx.make_node("Cast", [k_0d], attr={"to": onnx_pb.TensorProto.INT64})
+        k_1d = ctx.make_node("Unsqueeze", cast.output, attr={"axes": [0]})
+        ctx.replace_input(node, k_0d, k_1d.output[0])
+
 
 @tf_op("Tile")
 class Tile: