pnnx handle more softmin logsoftmax dtype, fuse static full range sli…

…ces to tensor_split
Tencent · Jan 2, 2024 · bed43b4 · bed43b4
1 parent 1e88fb8
commit bed43b4
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 3 deletions.
diff --git a/tools/pnnx/src/pass_level2/F_log_softmax.cpp b/tools/pnnx/src/pass_level2/F_log_softmax.cpp
@@ -25,7 +25,7 @@ class F_log_softmax : public GraphRewriterPass
 5 4
 pnnx.Input              input_0     0 1 input
 pnnx.Input              input_1     0 1 dim
-prim::Constant          op_0        0 1 dtype value=None
+prim::Constant          op_0        0 1 dtype value=*
 aten::log_softmax       op_1        3 1 input dim dtype out
 pnnx.Output             output      1 0 out
 )PNNXIR";

diff --git a/tools/pnnx/src/pass_level2/F_softmin.cpp b/tools/pnnx/src/pass_level2/F_softmin.cpp
@@ -26,7 +26,7 @@ class F_softmin : public GraphRewriterPass
 pnnx.Input              input_0     0 1 input
 pnnx.Input              input_1     0 1 dim
 aten::neg               op_0        1 1 input 6
-prim::Constant          op_1        0 1 dtype value=None
+prim::Constant          op_1        0 1 dtype value=*
 aten::softmax           op_2        3 1 6 dim dtype out
 pnnx.Output             output      1 0 out
 )PNNXIR";

diff --git a/tools/pnnx/src/pass_level5/fuse_slice_to_tensor_split.cpp b/tools/pnnx/src/pass_level5/fuse_slice_to_tensor_split.cpp
@@ -109,6 +109,12 @@ void fuse_slice_to_tensor_split(Graph& graph)
                     full_dimsize_slice = true;
                     break;
                 }
+                if (!op_in->shape.empty() && end2 == op_in->shape[dim])
+                {
+                    slice_n_ops.push_back(op2);
+                    full_dimsize_slice = true;
+                    break;
+                }
 
                 tensor_split_indices.push_back(end2);
                 slice_n_ops.push_back(op2);

diff --git a/tools/pnnx/tests/test_pnnx_fuse_slice_to_tensor_split.py b/tools/pnnx/tests/test_pnnx_fuse_slice_to_tensor_split.py
@@ -24,7 +24,7 @@ def forward(self, x, y, z):
         x0 = x[:3]
         x1 = x[3:]
 
-        z3 = z[:,:,7:]
+        z3 = z[:,:,7:z.size(2)]
         z2 = z[:,:,4:7]
 
         y0 = y[:2,:]