pnnx ncnn handles squeeze/unsqueeze the batch index (#6232)

Author: nihui

tools/pnnx/src/pass_ncnn/solve_batch_index.cpp

@@ -389,14 +389,14 @@ static void solve_batch_index_forward(Operand* operand)
                 dim += input_rank0;
 
             int batch_index_squeezed = batch_index;
-            if (dim >= 0 && dim < batch_index)
-            {
-                batch_index_squeezed = batch_index - 1;
-            }
             if (dim >= 0 && dim == batch_index)
             {
                 batch_index_squeezed = 233;
             }
+            else if (dim >= 0 && dim < batch_index)
+            {
+                batch_index_squeezed = batch_index - 1;
+            }
 
             Operand* r = op->outputs[0];
             if (r->params.find("__batch_index") == r->params.end())
@@ -413,6 +413,12 @@ static void solve_batch_index_forward(Operand* operand)
             if (dim < 0)
                 dim += input_rank0;
 
+            if (batch_index == 233)
+            {
+                // give up
+                return;
+            }
+
             int batch_index_unsqueezed = batch_index;
             if (dim >= 0 && dim <= batch_index)
             {
@@ -703,6 +709,12 @@ static void solve_batch_index_backward(Operand* operand)
         if (dim < 0)
             dim += input_rank0;
 
+        if (batch_index == 233)
+        {
+            // give up
+            return;
+        }
+
         int batch_index_unsqueezed = batch_index;
         if (dim >= 0 && dim <= batch_index)
         {
@@ -725,7 +737,11 @@ static void solve_batch_index_backward(Operand* operand)
             dim += input_rank0;
 
         int batch_index_squeezed = batch_index;
-        if (dim >= 0 && dim <= batch_index)
+        if (dim >= 0 && dim == batch_index)
+        {
+            batch_index_squeezed = 233;
+        }
+        else if (dim >= 0 && dim <= batch_index)
         {
             batch_index_squeezed = batch_index - 1;
         }