asserts LtNodes of SumNode with MulNode of Nodes (#2465)

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)

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