nebius/SWE-bench-extra huggingface dataset:
Repository: y0-causal-inference/y0

Issue:
Canonicalize does not work on Sum
```python
from y0.dsl import P, Sum, X, Y, Z, Product
from y0.mutate import canonicalize
expected = Sum[Z](P(Y,Z))
actual = Sum(P(Z, Y),[Z])
expected_vars = expected.get_variables()
ordering = list(expected_vars)
expected_canonical = canonicalize(expected, ordering)
actual_canonical = canonicalize(actual, ordering)

print(f"Expected: {expected_canonical}\nActual: {actual_canonical}")
```

```
Expected: [ sum_{Z} P(Y,Z) ]
Actual: [ sum_{Z} P(Z,Y) ]
```


Comments:

Patch:
diff --git a/src/y0/mutate/canonicalize_expr.py b/src/y0/mutate/canonicalize_expr.py
index 61924c0..b5122f7 100644
--- a/src/y0/mutate/canonicalize_expr.py
+++ b/src/y0/mutate/canonicalize_expr.py
@@ -64,9 +64,6 @@ class Canonicalizer:
         if isinstance(expression, Probability):  # atomic
             return self._canonicalize_probability(expression)
         elif isinstance(expression, Sum):
-            if isinstance(expression.expression, Probability):  # also atomic
-                return expression
-
             return Sum(
                 expression=self.canonicalize(expression.expression),
                 ranges=expression.ranges,


Test patch:
diff --git a/tests/test_mutate/test_canonicalize.py b/tests/test_mutate/test_canonicalize.py
index 06ed32a..c76d878 100644
--- a/tests/test_mutate/test_canonicalize.py
+++ b/tests/test_mutate/test_canonicalize.py
@@ -108,13 +108,17 @@ class TestCanonicalize(unittest.TestCase):
             self.assert_canonicalize(P(A & B | C), P(c1 & c2 | C), [A, B, C])
             # Two conditions, C and D
             for p1, p2 in itt.permutations([C, D]):
-                self.assert_canonicalize(P(A & B | C | D), P(c1 & c2 | (p1, p2)), [A, B, C, D])
+                expected = P(A & B | C | D)
+                expression = P(c1 & c2 | (p1, p2))
+                ordering = [A, B, C, D]
+                self.assert_canonicalize(expected, expression, ordering)
+                self.assert_canonicalize(Sum(expected), Sum(expression), ordering)
 
         for c1, c2, c3 in itt.permutations([A, B, C]):
             self.assert_canonicalize(P(A, B, C), P(c1, c2, c3), [A, B, C])
             for p1, p2, p3 in itt.permutations([X, Y, Z]):
-                self.assert_canonicalize(
-                    P(A & B & C | (X, Y, Z)),
-                    P(c1 & c2 & c3 | (p1 & p2 & p3)),
-                    [A, B, C, X, Y, Z],
-                )
+                expected = P(A & B & C | (X, Y, Z))
+                expression = P(c1 & c2 & c3 | (p1 & p2 & p3))
+                ordering = [A, B, C, X, Y, Z]
+                self.assert_canonicalize(expected, expression, ordering)
+                self.assert_canonicalize(Sum(expected), Sum(expression), ordering)