Fix Canonical Simplifier (#5505)

diff --git a/src/arith/canonical_simplify.cc b/src/arith/canonical_simplify.cc
index 2bb0189..a10db7a 100644 (file)
--- a/src/arith/canonical_simplify.cc
+++ b/src/arith/canonical_simplify.cc
@@ -870,6 +870,7 @@ SplitModConst(SplitExpr lhs, int64_t cval, DivMode div_mode) {
           lhs->upper_factor != SplitExprNode::kPosInf) {
         auto updated = ToSplitExpr(this->VisitExpr(ModImpl(
             lhs->index, make_const(lhs.dtype(), new_upper_factor), div_mode)));
+        updated.CopyOnWrite()->scale = lhs->scale;
         // re-apply the lower_factor
         if (lhs->lower_factor != 1) {
           return SplitDivConst(updated, lhs->lower_factor, div_mode);

diff --git a/tests/python/unittest/test_arith_canonical_simplify.py b/tests/python/unittest/test_arith_canonical_simplify.py
index 0dcf1fb..1791522 100644 (file)
--- a/tests/python/unittest/test_arith_canonical_simplify.py
+++ b/tests/python/unittest/test_arith_canonical_simplify.py
@@ -126,7 +126,7 @@ def test_floormod_simplify():
     x, y = te.var("x"), te.var("y")
     ck.verify(flm(flm((x*4) + y  - 466036, 24528) - 24512,  16),
               flm((x*4) + y  + 12, 16))
-
+    ck.verify(flm(flm((x*4), 16), 8), flm(x, 2) * 4)
 
 
 def test_canonical_mixed():