DecomposeComplexOps: Use static shape if available

Xilinx · Jul 21, 2023 · d91b608 · d91b608
1 parent 2174481
commit d91b608
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 11 deletions.
diff --git a/e2e_testing/xfail_sets.py b/e2e_testing/xfail_sets.py
@@ -1267,16 +1267,6 @@
 }) - {
 ### Test failing in make_fx_tosa but not in tosa
 
-    # 'tosa.const' op failed to verify that all of {value, output} have same shape
-    "BatchNorm1DModule_basic",
-    "BatchNorm1DWith2DInputModule_basic",
-    "BatchNorm2DModule_basic",
-    "BatchNorm3DModule_basic",
-
-    # 'tensor.empty' op incorrect number of dynamic sizes, has 1, expected 0
-    "BatchNorm1DStaticShapeModule_basic",
-    "ResNet18StaticModule_basic",
-
     # Dynamic shape, has extra unsupported broadcast ops
     "Matmul_3d",
 

diff --git a/lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp b/lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp
@@ -2768,7 +2768,7 @@ class DecomposeAtenNativeBatchNormOp
         loc, ListType::get(IntType::get(context)), runningStatsShape);
 
     SmallVector<int64_t> runningStatsShapeInt(inputRank, 1);
-    runningStatsShapeInt[1] = kUnknownSize;
+    runningStatsShapeInt[1] = runningMean.getType().cast<BaseTensorType>().getSizes()[0];
     Type dtype = input.getType().cast<ValueTensorType>().getOptionalDtype();
     Type reshapeType = ValueTensorType::get(
         context, llvm::ArrayRef(runningStatsShapeInt), dtype);