Merge pull request #808 from wjj19950828/fixed_nonzero

Fixed nonzero op
PaddlePaddle · Jun 8, 2022 · 7101837 · 7101837
2 parents 2cb5ba2 + 06a3deb
commit 7101837
Showing 1 changed file with 7 additions and 24 deletions.
diff --git a/x2paddle/op_mapper/onnx2paddle/opset9/opset.py b/x2paddle/op_mapper/onnx2paddle/opset9/opset.py
@@ -1880,30 +1880,13 @@ def Where(self, node):
     @print_mapping_info
     def NonZero(self, node):
         val_x = self.graph.get_input_node(node, idx=0, copy=True)
-        val_x_dim = len(val_x.out_shapes[0])
-        if val_x_dim == 1:
-            self.paddle_graph.add_layer(
-                "paddle.nonzero",
-                inputs={"x": val_x.name},
-                outputs=[val_x.name])
-            self.paddle_graph.add_layer(
-                "paddle.transpose",
-                inputs={"x": val_x.name},
-                outputs=[node.layer_name],
-                perm=[1, 0])
-        if val_x_dim > 1:
-            self.paddle_graph.add_layer(
-                "paddle.nonzero",
-                inputs={"x": val_x.name},
-                outputs=[val_x.name])
-            self.paddle_graph.add_layer(
-                "paddle.split",
-                inputs={"x": val_x.name},
-                outputs=[val_x.name],
-                num_or_sections=1,
-                axis=val_x_dim)
-            self.paddle_graph.add_layer(
-                "paddle.concat", inputs={"x": val_x.name}, outputs=[node.name])
+        self.paddle_graph.add_layer(
+            "paddle.nonzero",
+            inputs={"x": val_x.name},
+            outputs=[val_x.name],
+            as_tuple=True)
+        self.paddle_graph.add_layer(
+            "paddle.concat", inputs={"x": val_x.name}, outputs=[node.name])
 
     @print_mapping_info
     def Identity(self, node):