@@ -18792,30 +18792,33 @@ struct ExtendSlice final
1879218792 int64_t limit_d = limits[d];
1879318793 int64_t size_d = originalShape[d];
1879418794
18795- // Calculate the parameters for the new slice operation on the original operand.
18796- // The new slice covers the part of the original tensor that is visible in the final output.
18795+ // Calculate the parameters for the new slice operation on the original
18796+ // operand. The new slice covers the part of the original tensor that is
18797+ // visible in the final output.
1879718798 new_starts[d] = std::max((int64_t)0, start_d - lhs);
1879818799 new_limits[d] = std::min(size_d, limit_d - lhs);
1879918800
1880018801 // Calculate the new padding amounts for the extend operation.
18801- // new_lhs is the size of the overlap between the slice and the prepended padding.
18802+ // new_lhs is the size of the overlap between the slice and the prepended
18803+ // padding.
1880218804 int64_t new_lhs = std::max((int64_t)0, std::min(limit_d, lhs) - start_d);
18803- // new_rhs is the size of the overlap between the slice and the appended padding.
18804- int64_t new_rhs = std::max((int64_t)0, limit_d - std::max(start_d, lhs + size_d));
18805+ // new_rhs is the size of the overlap between the slice and the appended
18806+ // padding.
18807+ int64_t new_rhs =
18808+ std::max((int64_t)0, limit_d - std::max(start_d, lhs + size_d));
1880518809
1880618810 // Create the new slice on the original tensor.
1880718811 auto newSlice = rewriter.create<stablehlo::SliceOp>(
1880818812 op.getLoc(), operand, new_starts, new_limits, new_strides);
1880918813
1881018814 // Create the new extend on the newly sliced tensor.
18811- rewriter.replaceOpWithNewOp<enzymexla::ExtendOp>(
18812- op, op.getType(), newSlice, new_lhs, new_rhs, d);
18815+ rewriter.replaceOpWithNewOp<enzymexla::ExtendOp>(op, op.getType(), newSlice,
18816+ new_lhs, new_rhs, d);
1881318817
1881418818 return success();
1881518819 }
1881618820};
1881718821
18818-
1881918822struct SliceExtend final
1882018823 : CheckedOpRewritePattern<enzymexla::ExtendOp, SliceExtend> {
1882118824 using CheckedOpRewritePattern::CheckedOpRewritePattern;
0 commit comments