Fix ReduceArg of TFLite importer

kendryte · Mar 6, 2025 · b39a7bd · b39a7bd
1 parent 0273c62
commit b39a7bd
Showing 1 changed file with 2 additions and 13 deletions.
diff --git a/src/Nncase.Importer/TFLite/Reduce.cs b/src/Nncase.Importer/TFLite/Reduce.cs
@@ -12,7 +12,7 @@ public partial class TFLiteImporter
         private Expr VisitReduce(in tflite.Operator op, ReduceOp reduceOp, float initValue)
         {
             var (input, axis) = GetInputExprs(op, 0, 1);
-            return Reduce(reduceOp, input, ProcAxis(axis), initValue, op.BuiltinOptionsAsReducerOptions().KeepDims);
+            return Reduce(reduceOp, input, axis, initValue, op.BuiltinOptionsAsReducerOptions().KeepDims);
         }
 
         private Expr VisitReduceArg(in tflite.Operator op, ReduceArgOp reduceArgOp)
@@ -25,18 +25,7 @@ private Expr VisitReduceArg(in tflite.Operator op, ReduceArgOp reduceArgOp)
                 _ => throw new ArgumentOutOfRangeException(nameof(reduceArgOp), reduceArgOp, null),
             };
 
-            return ReduceArg(reduceArgOp, (PrimType)GetDataType(outType), input, ProcAxis(axis), false, false);
-        }
-
-        private Expr ProcAxis(Expr axis)
-        {
-            if (axis is TensorConst axisValue)
-            {
-                // scalar to array
-                return axisValue.Value.ToArray<int>();
-            }
-
-            return axis;
+            return ReduceArg(reduceArgOp, (PrimType)GetDataType(outType), input, axis, false, false);
         }
     }
 }