Fix schedules

Author: saulshanabrook

python/egglog/egraph.py

@@ -1565,7 +1565,7 @@ def __add__(self, other: Schedule) -> Schedule:
         """
         Run two schedules in sequence.
         """
-        return Schedule(Thunk.fn(Declarations.create, self, other), SequenceDecl((self.schedule, other.schedule)))
+        return Schedule(partial(Declarations.create, self, other), SequenceDecl((self.schedule, other.schedule)))
 
 
 @dataclass
@@ -2046,7 +2046,7 @@ def run(ruleset: Ruleset | None = None, *until: FactLike, scheduler: BackOff | N
     """
     facts = _fact_likes(until)
     return Schedule(
-        Thunk.fn(Declarations.create, ruleset, *facts),
+        partial(Declarations.create, ruleset, *facts),
         RunDecl(
             ruleset.__egg_ident__ if ruleset else Ident(""),
             tuple(f.fact for f in facts),
@@ -2089,7 +2089,7 @@ def seq(*schedules: Schedule) -> Schedule:
     """
     Run a sequence of schedules.
     """
-    return Schedule(Thunk.fn(Declarations.create, *schedules), SequenceDecl(tuple(s.schedule for s in schedules)))
+    return Schedule(partial(Declarations.create, *schedules), SequenceDecl(tuple(s.schedule for s in schedules)))
 
 
 def _action_likes(action_likes: Iterable[ActionLike]) -> tuple[Action, ...]:

python/tests/test_array_api.py

@@ -172,6 +172,18 @@ def test_normalize_reshape_shape(self):
             array_api_schedule,
         )
 
+    def test_reshape_after_schedule_decls_access(self):
+        _ = array_api_schedule.__egg_decls__
+
+        x = NDArray.var("x")
+        assume_shape(x, TupleInt((Int(5),)))
+        res = reshape(x, TupleInt((-1,)))
+        egraph = EGraph()
+        egraph.register(res)
+        egraph.run(array_api_schedule)
+
+        egraph.check(eq(res).to(x))
+
     def test_reshape_vec_noop(self):
         x = NDArray.var("x")
         assume_shape(x, TupleInt((Int(5),)))

python/tests/test_high_level.py

@@ -1117,6 +1117,20 @@ def __sub__(self, other: E) -> E: ...
 
 
 class TestScheduler:
+    def test_seq_schedule_decls_track_ruleset_updates(self):
+        egraph = EGraph()
+
+        rel = relation("rel_live", i64)
+        live_rules = ruleset(name="live-rules")
+        schedule = seq(live_rules, run()).saturate()
+        _ = schedule.__egg_decls__
+
+        live_rules.register(rule(rel(i64(0))).then(rel(i64(1))))
+
+        egraph.register(rel(i64(0)))
+        egraph.run(schedule)
+        egraph.check(rel(i64(1)))
+
     def test_sequence_repeat_saturate(self):
         """
         Mirrors the scheduling example: alternate step-right and step-left,