pnnx match onnx zeros ones (#5832)

Tencent · Dec 19, 2024 · 57fac4a · 57fac4a
1 parent ad6d84e
commit 57fac4a
Show file tree

Hide file tree

Showing 3 changed files with 53 additions and 1 deletion.
diff --git a/tools/pnnx/src/pass_level2/torch_full.cpp b/tools/pnnx/src/pass_level2/torch_full.cpp
@@ -91,6 +91,6 @@ pnnx.Output             output      1 0 out
     }
 };
 
-REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_full_onnx, 20)
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_full_onnx, 21)
 
 } // namespace pnnx
diff --git a/tools/pnnx/src/pass_level2/torch_ones.cpp b/tools/pnnx/src/pass_level2/torch_ones.cpp
@@ -65,4 +65,30 @@ pnnx.Output             output      1 0 out
 
 REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_ones, 20)
 
+class torch_ones_onnx : public GraphRewriterPass
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+3 2
+pnnx.Input              input       0 1 size
+ConstantOfShape         op_0        1 1 size out value=1.0
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    const char* type_str() const
+    {
+        return "torch.ones";
+    }
+
+    void write(Operator* op, const std::map<std::string, Parameter>& /*captured_params*/) const
+    {
+        op->params["dtype"] = Parameter();
+    }
+};
+
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_ones_onnx, 20)
+
 } // namespace pnnx
diff --git a/tools/pnnx/src/pass_level2/torch_zeros.cpp b/tools/pnnx/src/pass_level2/torch_zeros.cpp
@@ -65,4 +65,30 @@ pnnx.Output             output      1 0 out
 
 REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_zeros, 20)
 
+class torch_zeros_onnx : public GraphRewriterPass
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+3 2
+pnnx.Input              input       0 1 size
+ConstantOfShape         op_0        1 1 size out value=0.0
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    const char* type_str() const
+    {
+        return "torch.zeros";
+    }
+
+    void write(Operator* op, const std::map<std::string, Parameter>& /*captured_params*/) const
+    {
+        op->params["dtype"] = Parameter();
+    }
+};
+
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_zeros_onnx, 20)
+
 } // namespace pnnx