Merge pull request #1117 from xadupre/ip

Fixes #595, support operator InvertPermutation

Author: xadupre

tests/test_backend.py

@@ -3676,6 +3676,19 @@ def func3(x):
         with self.assertRaises(ValueError):
             self._run_test_case(func3, [_OUTPUT], {_INPUT: x_val})
 
+    @check_opset_min_version(11, "topk")
+    def test_invert_permutation(self):
+
+        def func(x):
+            op_ = tf.math.invert_permutation(x)
+            return tf.identity(op_, name=_TFOUTPUT)
+
+        x_val = np.array([0, 1, 2, 3], dtype=np.int64)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
+
+        x_val = np.array([1, 5, 2, 0, 3, 4], dtype=np.int64)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
+
 
 if __name__ == '__main__':
     unittest_main()

tf2onnx/onnx_opset/math.py

@@ -700,3 +700,34 @@ def atan2(y, x):
             op_name_scope=node.name + 'all',
             shapes=[shape], dtypes=[onnx_dtype])
         ctx.replace_all_inputs(node.output[0], last_node.output[0])  # ops=ctx.get_nodes()
+
+
+@tf_op("InvertPermutation")
+class InvertPermutationOp:
+
+    @classmethod
+    def version_11(cls, ctx, node, **kwargs):
+
+        supported_dtypes = [onnx_pb.TensorProto.INT32, onnx_pb.TensorProto.INT64]
+        onnx_dtype = ctx.get_dtype(node.input[0])
+        utils.make_sure(onnx_dtype in supported_dtypes, "InvertPermutation only applies on INT32, INT64.")
+
+        shape = ctx.get_shape(node.input[0])
+
+        shape_node = ctx.make_node(
+            "Shape", inputs=node.input, name=utils.make_name(node.name + '_shape'))
+
+        neg_node = ctx.make_node(
+            "Neg", inputs=node.input, name=utils.make_name(node.name + '_neg'))
+
+        topk_node = ctx.make_node(
+            "TopK", inputs=[neg_node.output[0], shape_node.output[0]],
+            name=utils.make_name(node.name + '_topk'), output_count=2)
+
+        ctx.remove_node(node.name)
+
+        last_node = ctx.make_node(
+            "Identity", inputs=topk_node.output[1:], name=utils.make_name(node.name + '_indices'),
+            shapes=[shape], dtypes=[onnx_dtype])
+
+        ctx.replace_all_inputs(node.output[0], last_node.output[0])  # ops=ctx.get_nodes()