Fix AtenArangeStartStepOp dynamic end support

llvm · Nov 20, 2024 · f7cd0fc · f7cd0fc
1 parent 35e20e0
commit f7cd0fc
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp b/lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp
@@ -5604,11 +5604,13 @@ class DecomposeAtenNonzeroOp : public OpRewritePattern<AtenNonzeroOp> {
                                                     subtracted, c(0));
 
     //     iota = torch.tensor(range(len(t))) * nonzero_mask.int()
+    Value constZero = rewriter.create<Torch::ConstantIntOp>(
+        loc, rewriter.getI64IntegerAttr(0));
+    Value end =
+        rewriter.create<AtenSizeIntOp>(loc, flattenedInput, /*dim=*/constZero);
     Value rangeTensor = rewriter.create<AtenArangeStartStepOp>(
-        loc, cumulativeSumType, c(0),
-        rewriter.create<ConstantIntOp>(
-            loc, rewriter.getI64IntegerAttr(flattenedInputType.getSizes()[0])),
-        one, noneCst, noneCst, noneCst, noneCst);
+        loc, cumulativeSumType, c(0), end, one, noneCst, noneCst, noneCst,
+        noneCst);
     Value multiplied = rewriter.create<AtenMulTensorOp>(loc, cumulativeSumType,
                                                         rangeTensor, intMask);