Fix Canonical Simplifier (apache#5505)

neo-ai · Jun 18, 2020 · ae62e23 · ae62e23
1 parent 6488bbd
commit ae62e23
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git 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
@@ -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():