Fixes #595, support operator InvertPermutation

Signed-off-by: xavier dupré <[email protected]>

Author: sdpython

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(10, "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

@@ -697,3 +697,38 @@ 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:
+    supported_dtypes = [
+        onnx_pb.TensorProto.INT32,
+        onnx_pb.TensorProto.INT64,
+    ]
+
+    @classmethod
+    def version_10(cls, ctx, node, **kwargs):
+
+        onnx_dtype = ctx.get_dtype(node.input[0])
+        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_unused = utils.make_name(node.name + '_topk')
+        topk_indices = utils.make_name(node.name + '_indices')
+        outputs = [topk_unused, utils.port_name(topk_indices, 1)]
+        topk_node = ctx.make_node(
+            "TopK", inputs=[neg_node.output[0], shape_node.output[0]],
+            name=utils.make_name(node.name + '_topk'), outputs=outputs)
+
+        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()