[Torch Dialect] Fix small bugs in decompose-complex-ops pass, e.g. mi…

…ssing return sentence (#2409)
llvm · Aug 22, 2023 · b552d4e · b552d4e
1 parent 4c9d234
commit b552d4e
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp b/lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp
@@ -1292,8 +1292,8 @@ class DecomposeAtenRepeatOp : public OpRewritePattern<AtenRepeatOp> {
 
     SmallVector<Value> unsqueezedSizes, expandedSizes, reshapedSizes;
     SmallVector<int64_t> unsqueezedIntSizes, expandedIntSizes;
+    assert(repeats.size() >= rank && "leadingRank should greater than 0");
     auto leadingRank = repeats.size() - rank;
-    assert(leadingRank >= 0 && "leadingRank should greater than 0");
     for (size_t i = 0; i < leadingRank; ++i) {
       insertDimSizes(unsqueezedSizes, unsqueezedIntSizes, ArrayRef<Value>{one});
       insertDimSizes(expandedSizes, expandedIntSizes,
@@ -4808,7 +4808,7 @@ class DecomposeAtenIndexTensorOp : public OpRewritePattern<AtenIndexTensorOp> {
     int64_t inputRank = inputSizes.size();
     auto outputType = op.getType().cast<BaseTensorType>();
     if (!outputType.hasSizes()) {
-      rewriter.notifyMatchFailure(
+      return rewriter.notifyMatchFailure(
           op, "only output with shape information is supported");
     }
     auto outputRank = outputType.getSizes().size();