diff --git a/test/unit/test_symbolic.py b/test/unit/test_symbolic.py index 83a3460085..673657d1d1 100644 --- a/test/unit/test_symbolic.py +++ b/test/unit/test_symbolic.py @@ -399,8 +399,8 @@ class TestSymbolicSymbolicOps(unittest.TestCase): b = Variable("b", 1, 2) c = Variable("c", 1, 2) x = SumNode([MulNode(a, b), c]) - assert isinstance((x < 3), Node) and (x < 3) == 0 - assert isinstance((x < 4), LtNode) and (x < 4).min == 0 and (x < 4).max == 1 + with self.assertRaises(AssertionError): + lt3 = (x < 3) def test_num_node_mul_node(self): a = Variable("a", 1, 5) diff --git a/tinygrad/shape/symbolic.py b/tinygrad/shape/symbolic.py index d3b84cce6b..87062e012f 100644 --- a/tinygrad/shape/symbolic.py +++ b/tinygrad/shape/symbolic.py @@ -286,11 +286,12 @@ class SumNode(RedNode): else: new_sum.append(x) lhs = Node.sum(new_sum) nodes = lhs.nodes if isinstance(lhs, SumNode) else [lhs] + assert all(not isinstance(node, MulNode) or isinstance(node.b, int) for node in nodes), "not supported" muls, others = partition(nodes, lambda x: isinstance(x, MulNode) and x.b > 0 and x.max >= b) if muls: # NOTE: gcd in python 3.8 takes exactly 2 args mul_gcd = b - for x in muls: mul_gcd = gcd(mul_gcd, x.b) if isinstance(x.b, int) else 1 + for x in muls: mul_gcd = gcd(mul_gcd, x.b) # type: ignore # mypy cannot tell that x.b is int here due to assert above all_others = Variable.sum(others) if all_others.min >= 0 and all_others.max < mul_gcd: lhs, b = Variable.sum([mul//mul_gcd for mul in muls]), b//mul_gcd