egglog 9.0.1__cp311-cp311-win_amd64.whl → 10.0.0__cp311-cp311-win_amd64.whl

egglog/__init__.py

@@ -3,6 +3,7 @@ Package for creating e-graphs in Python.
 """
 
 from . import config, ipython_magic  # noqa: F401
+from .bindings import EggSmolError  # noqa: F401
 from .builtins import *  # noqa: UP029
 from .conversion import ConvertError, convert, converter, get_type_args  # noqa: F401
 from .egraph import *

egglog/builtins.py

@@ -12,11 +12,9 @@ from typing import TYPE_CHECKING, Generic, Protocol, TypeAlias, TypeVar, Union,
 
 from typing_extensions import TypeVarTuple, Unpack
 
-from . import bindings
 from .conversion import convert, converter, get_type_args
 from .declarations import *
-from .egraph import BaseExpr, BuiltinExpr, EGraph, expr_fact, function, get_current_ruleset, method
-from .egraph_state import GLOBAL_PY_OBJECT_SORT
+from .egraph import BaseExpr, BuiltinExpr, expr_fact, function, get_current_ruleset, method
 from .functionalize import functionalize
 from .runtime import RuntimeClass, RuntimeExpr, RuntimeFunction
 from .thunk import Thunk
@@ -32,6 +30,7 @@ __all__ = [
     "BigRatLike",
     "Bool",
     "BoolLike",
+    "BuiltinEvalError",
     "Map",
     "MapLike",
     "MultiSet",
@@ -56,6 +55,17 @@ __all__ = [
 ]
 
 
+class BuiltinEvalError(Exception):
+    """
+    Raised when an builtin cannot be evaluated into a Python primitive because it is complex.
+
+    Try extracting this expression first.
+    """
+
+    def __str__(self) -> str:
+        return f"Cannot evaluate builtin expression into a Python primitive. Try extracting this expression first: {super().__str__()}"
+
+
 class Unit(BuiltinExpr, egg_sort="Unit"):
     """
     The unit type. This is used to reprsent if a value exists in the e-graph or not.
@@ -72,8 +82,8 @@ class String(BuiltinExpr):
     @method(preserve=True)
     def eval(self) -> str:
         value = _extract_lit(self)
-        assert isinstance(value, bindings.String)
-        return value.value
+        assert isinstance(value, str)
+        return value
 
     def __init__(self, value: str) -> None: ...
 
@@ -97,8 +107,8 @@ class Bool(BuiltinExpr, egg_sort="bool"):
     @method(preserve=True)
     def eval(self) -> bool:
         value = _extract_lit(self)
-        assert isinstance(value, bindings.Bool)
-        return value.value
+        assert isinstance(value, bool)
+        return value
 
     @method(preserve=True)
     def __bool__(self) -> bool:
@@ -132,8 +142,8 @@ class i64(BuiltinExpr):  # noqa: N801
     @method(preserve=True)
     def eval(self) -> int:
         value = _extract_lit(self)
-        assert isinstance(value, bindings.Int)
-        return value.value
+        assert isinstance(value, int)
+        return value
 
     @method(preserve=True)
     def __index__(self) -> int:
@@ -251,8 +261,8 @@ class f64(BuiltinExpr):  # noqa: N801
     @method(preserve=True)
     def eval(self) -> float:
         value = _extract_lit(self)
-        assert isinstance(value, bindings.Float)
-        return value.value
+        assert isinstance(value, float)
+        return value
 
     @method(preserve=True)
     def __float__(self) -> float:
@@ -340,9 +350,12 @@ class Map(BuiltinExpr, Generic[T, V]):
         expr = cast("RuntimeExpr", self)
         d = {}
         while call.callable != ClassMethodRef("Map", "empty"):
-            assert call.callable == MethodRef("Map", "insert")
+            msg = "Map can only be evaluated if it is empty or a series of inserts."
+            if call.callable != MethodRef("Map", "insert"):
+                raise BuiltinEvalError(msg)
             call_typed, k_typed, v_typed = call.args
-            assert isinstance(call_typed.expr, CallDecl)
+            if not isinstance(call_typed.expr, CallDecl):
+                raise BuiltinEvalError(msg)
             k = cast("T", expr.__with_expr__(k_typed))
             v = cast("V", expr.__with_expr__(v_typed))
             d[k] = v
@@ -404,7 +417,9 @@ class Set(BuiltinExpr, Generic[T]):
     @method(preserve=True)
     def eval(self) -> set[T]:
         call = _extract_call(self)
-        assert call.callable == InitRef("Set")
+        if call.callable != InitRef("Set"):
+            msg = "Set can only be initialized with the Set constructor."
+            raise BuiltinEvalError(msg)
         return {cast("T", cast("RuntimeExpr", self).__with_expr__(x)) for x in call.args}
 
     @method(preserve=True)
@@ -466,7 +481,9 @@ class MultiSet(BuiltinExpr, Generic[T]):
     @method(preserve=True)
     def eval(self) -> list[T]:
         call = _extract_call(self)
-        assert call.callable == InitRef("MultiSet")
+        if call.callable != InitRef("MultiSet"):
+            msg = "MultiSet can only be initialized with the MultiSet constructor."
+            raise BuiltinEvalError(msg)
         return [cast("T", cast("RuntimeExpr", self).__with_expr__(x)) for x in call.args]
 
     @method(preserve=True)
@@ -513,11 +530,15 @@ class Rational(BuiltinExpr):
     @method(preserve=True)
     def eval(self) -> Fraction:
         call = _extract_call(self)
-        assert call.callable == InitRef("Rational")
+        if call.callable != InitRef("Rational"):
+            msg = "Rational can only be initialized with the Rational constructor."
+            raise BuiltinEvalError(msg)
 
         def _to_int(e: TypedExprDecl) -> int:
             expr = e.expr
-            assert isinstance(expr, LitDecl)
+            if not isinstance(expr, LitDecl):
+                msg = "Rational can only be initialized with literals"
+                raise BuiltinEvalError(msg)
             assert isinstance(expr.value, int)
             return expr.value
 
@@ -596,9 +617,13 @@ class BigInt(BuiltinExpr):
     @method(preserve=True)
     def eval(self) -> int:
         call = _extract_call(self)
-        assert call.callable == ClassMethodRef("BigInt", "from_string")
+        if call.callable != ClassMethodRef("BigInt", "from_string"):
+            msg = "BigInt can only be initialized with the BigInt constructor."
+            raise BuiltinEvalError(msg)
         (s,) = call.args
-        assert isinstance(s.expr, LitDecl)
+        if not isinstance(s.expr, LitDecl):
+            msg = "BigInt can only be initialized with literals"
+            raise BuiltinEvalError(msg)
         assert isinstance(s.expr.value, str)
         return int(s.expr.value)
 
@@ -717,14 +742,19 @@ class BigRat(BuiltinExpr):
     @method(preserve=True)
     def eval(self) -> Fraction:
         call = _extract_call(self)
-        assert call.callable == InitRef("BigRat")
+        if call.callable != InitRef("BigRat"):
+            msg = "BigRat can only be initialized with the BigRat constructor."
+            raise BuiltinEvalError(msg)
 
         def _to_fraction(e: TypedExprDecl) -> Fraction:
             expr = e.expr
-            assert isinstance(expr, CallDecl)
-            assert expr.callable == ClassMethodRef("BigInt", "from_string")
+            if not isinstance(expr, CallDecl) or expr.callable != ClassMethodRef("BigInt", "from_string"):
+                msg = "BigRat can only be initialized BigInt strings"
+                raise BuiltinEvalError(msg)
             (s,) = expr.args
-            assert isinstance(s.expr, LitDecl)
+            if not isinstance(s.expr, LitDecl):
+                msg = "BigInt can only be initialized with literals"
+                raise BuiltinEvalError(msg)
             assert isinstance(s.expr.value, str)
             return Fraction(s.expr.value)
 
@@ -821,7 +851,10 @@ class Vec(BuiltinExpr, Generic[T]):
         call = _extract_call(self)
         if call.callable == ClassMethodRef("Vec", "empty"):
             return ()
-        assert call.callable == InitRef("Vec")
+
+        if call.callable != InitRef("Vec"):
+            msg = "Vec can only be initialized with the Vec constructor."
+            raise BuiltinEvalError(msg)
         return tuple(cast("T", cast("RuntimeExpr", self).__with_expr__(x)) for x in call.args)
 
     @method(preserve=True)
@@ -889,10 +922,11 @@ VecLike: TypeAlias = Vec[T] | tuple[TO, ...] | list[TO]
 class PyObject(BuiltinExpr):
     @method(preserve=True)
     def eval(self) -> object:
-        report = (EGraph.current or EGraph())._run_extract(cast("RuntimeExpr", self), 0)
-        assert isinstance(report, bindings.Best)
-        expr = report.termdag.term_to_expr(report.term, bindings.PanicSpan())
-        return GLOBAL_PY_OBJECT_SORT.load(expr)
+        expr = cast("RuntimeExpr", self).__egg_typed_expr__.expr
+        if not isinstance(expr, PyObjectDecl):
+            msg = "PyObject can only be evaluated if it is a PyObject literal"
+            raise BuiltinEvalError(msg)
+        return expr.value
 
     def __init__(self, value: object) -> None: ...
 
@@ -1027,22 +1061,23 @@ def value_to_annotation(a: object) -> type | None:
 converter(FunctionType, UnstableFn, _convert_function)
 
 
-def _extract_lit(e: BaseExpr) -> bindings._Literal:
+def _extract_lit(e: BaseExpr) -> LitType:
     """
     Special case extracting literals to make this faster by using termdag directly.
     """
-    report = (EGraph.current or EGraph())._run_extract(cast("RuntimeExpr", e), 0)
-    assert isinstance(report, bindings.Best)
-    term = report.term
-    assert isinstance(term, bindings.TermLit)
-    return term.value
+    expr = cast("RuntimeExpr", e).__egg_typed_expr__.expr
+    if not isinstance(expr, LitDecl):
+        msg = "Expected a literal"
+        raise BuiltinEvalError(msg)
+    return expr.value
 
 
 def _extract_call(e: BaseExpr) -> CallDecl:
     """
     Extracts the call form of an expression
     """
-    extracted = cast("RuntimeExpr", (EGraph.current or EGraph()).extract(e))
-    expr = extracted.__egg_typed_expr__.expr
-    assert isinstance(expr, CallDecl)
+    expr = cast("RuntimeExpr", e).__egg_typed_expr__.expr
+    if not isinstance(expr, CallDecl):
+        msg = "Expected a call expression"
+        raise BuiltinEvalError(msg)
     return expr

egglog/conversion.py

@@ -149,6 +149,12 @@ def min_convertable_tp(a: object, b: object, name: str) -> JustTypeRef:
     decls = _retrieve_conversion_decls()
     a_tp = _get_tp(a)
     b_tp = _get_tp(b)
+    # Make sure at least one of the types has the method, to avoid issue with == upcasting improperly
+    if not (
+        (isinstance(a_tp, JustTypeRef) and decls.has_method(a_tp.name, name))
+        or (isinstance(b_tp, JustTypeRef) and decls.has_method(b_tp.name, name))
+    ):
+        raise ConvertError(f"Neither {a_tp} nor {b_tp} has method {name}")
     a_converts_to = {
         to: c for ((from_, to), (c, _)) in CONVERSIONS.items() if from_ == a_tp and decls.has_method(to.name, name)
     }

egglog/egraph.py

@@ -4,7 +4,7 @@ import contextlib
 import inspect
 import pathlib
 import tempfile
-from collections.abc import Callable, Generator, Iterable, Iterator
+from collections.abc import Callable, Generator, Iterable
 from contextvars import ContextVar
 from dataclasses import InitVar, dataclass, field
 from functools import partial
@@ -17,7 +17,6 @@ from typing import (
     Generic,
     Literal,
     Never,
-    Protocol,
     TypeAlias,
     TypedDict,
     TypeVar,
@@ -85,7 +84,6 @@ __all__ = [
     "set_",
     "simplify",
     "subsume",
-    "try_evaling",
     "union",
     "unstable_combine_rulesets",
     "var",
@@ -847,11 +845,10 @@ class EGraph:
     Can run actions, check facts, run schedules, or extract minimal cost expressions.
     """
 
-    current: ClassVar[EGraph | None] = None
     seminaive: InitVar[bool] = True
     save_egglog_string: InitVar[bool] = False
 
-    _state: EGraphState = field(init=False)
+    _state: EGraphState = field(init=False, repr=False)
     # For pushing/popping with egglog
     _state_stack: list[EGraphState] = field(default_factory=list, repr=False)
     # For storing the global "current" egraph
@@ -1200,16 +1197,6 @@ class EGraph:
             if visualize:
                 VisualizerWidget(egraphs=egraphs).display_or_open()
 
-    @contextlib.contextmanager
-    def set_current(self) -> Iterator[None]:
-        """
-        Context manager that will set the current egraph. It will be set back after.
-        """
-        prev_current = EGraph.current
-        EGraph.current = self
-        yield
-        EGraph.current = prev_current
-
     @property
     def _egraph(self) -> bindings.EGraph:
         return self._state.egraph
@@ -1303,8 +1290,6 @@ class Schedule(DelayedDeclerations):
     A composition of some rulesets, either composing them sequentially, running them repeatedly, running them till saturation, or running until some facts are met
     """
 
-    current: ClassVar[Schedule | None] = None
-
     # Defer declerations so that we can have rule generators that used not yet defined yet
     schedule: ScheduleDecl
 
@@ -1332,16 +1317,6 @@ class Schedule(DelayedDeclerations):
         """
         return Schedule(Thunk.fn(Declarations.create, self, other), SequenceDecl((self.schedule, other.schedule)))
 
-    @contextlib.contextmanager
-    def set_current(self) -> Iterator[None]:
-        """
-        Context manager that will set the current schedule. It will be set back after
-        """
-        prev_current = Schedule.current
-        Schedule.current = self
-        yield
-        Schedule.current = prev_current
-
 
 @dataclass
 class Ruleset(Schedule):
@@ -1488,9 +1463,12 @@ class Fact:
 
     def __bool__(self) -> bool:
         """
-        Returns True if the two sides of an equality are equal in the egraph or the expression is in the egraph.
+        Returns True if the two sides of an equality are structurally equal.
         """
-        return (EGraph.current or EGraph()).check_bool(self)
+        if not isinstance(self.fact, EqDecl):
+            msg = "Can only check equality facts"
+            raise TypeError(msg)
+        return self.fact.left == self.fact.right
 
 
 @dataclass
@@ -1876,34 +1854,3 @@ def set_current_ruleset(r: Ruleset | None) -> Generator[None, None, None]:
         yield
     finally:
         _CURRENT_RULESET.reset(token)
-
-
-T_co = TypeVar("T_co", covariant=True)
-
-
-class _EvalsTo(Protocol, Generic[T_co]):
-    def eval(self) -> T_co: ...
-
-
-def try_evaling(schedule: Schedule, expr: Expr, prim_expr: _EvalsTo[T]) -> T:
-    """
-    Try evaling the expression that will result in a primitive expression being fill.
-    if it fails, display the egraph and raise an error.
-    """
-    egraph = EGraph.current or EGraph()
-    with egraph.set_current():
-        try:
-            return prim_expr.eval()
-        except BaseException:  # noqa: S110
-            pass
-    # If this primitive doesn't exist in the egraph, we need to try to create it by
-    # registering the expression and running the schedule
-    egraph.register(expr)
-    egraph.run(Schedule.current or schedule)
-    try:
-        with egraph.set_current():
-            return prim_expr.eval()
-    except BaseException as e:
-        # egraph.display(n_inline_leaves=1, split_primitive_outputs=True)
-        e.add_note(f"Cannot evaluate {egraph.extract(expr)}")
-        raise

egglog/examples/bignum.py

@@ -12,14 +12,15 @@ x = BigInt(-1234)
 y = BigInt.from_string("2")
 z = BigRat(x, y)
 
-assert z.numer.to_string() == "-617"
+egraph = EGraph()
+
+assert egraph.extract(z.numer.to_string()).eval() == "-617"
 
 
 @function
 def bignums(x: BigInt, y: BigInt) -> BigRat: ...
 
 
-egraph = EGraph()
 egraph.register(set_(bignums(x, y)).to(z))
 
 c = var("c", BigRat)

egglog/examples/eqsat_basic.py

@@ -28,18 +28,17 @@ expr2 = Num(6) + Num(2) * Num.var("x")
 a, b, c = vars_("a b c", Num)
 i, j = vars_("i j", i64)
 
-check(
-    # Check that these expressions are equal
-    eq(expr1).to(expr2),
-    # After running these rules, up to ten times
+egraph = EGraph()
+egraph.register(expr1, expr2)
+
+egraph.run(
     ruleset(
         rewrite(a + b).to(b + a),
         rewrite(a * (b + c)).to((a * b) + (a * c)),
         rewrite(Num(i) + Num(j)).to(Num(i + j)),
         rewrite(Num(i) * Num(j)).to(Num(i * j)),
     )
-    * 10,
-    # On these two initial expressions
-    expr1,
-    expr2,
+    * 10
 )
+
+egraph.check(expr1 == expr2)

egglog/examples/multiset.py

@@ -29,33 +29,32 @@ egraph = EGraph()
 xs = MultiSet(Math(1), Math(2), Math(3))
 egraph.register(xs)
 
-with egraph.set_current():
-    assert xs == MultiSet(Math(1), Math(3), Math(2))
-    assert xs != MultiSet(Math(1), Math(1), Math(2), Math(3))
+egraph.check(xs == MultiSet(Math(1), Math(3), Math(2)))
+egraph.check_fail(xs == MultiSet(Math(1), Math(1), Math(2), Math(3)))
 
-    assert Counter(xs) == Counter({Math(1): 1, Math(2): 1, Math(3): 1})
+assert Counter(egraph.extract(xs).eval()) == Counter({Math(1): 1, Math(2): 1, Math(3): 1})
 
-    inserted = MultiSet(Math(1), Math(2), Math(3), Math(4))
-    egraph.register(inserted)
-    assert xs.insert(Math(4)) == inserted
 
-    assert xs.contains(Math(1))
-    assert xs.not_contains(Math(4))
-    assert Math(1) in xs
-    assert Math(4) not in xs
+inserted = MultiSet(Math(1), Math(2), Math(3), Math(4))
+egraph.register(inserted)
+egraph.check(xs.insert(Math(4)) == inserted)
+egraph.check(xs.contains(Math(1)))
+egraph.check(xs.not_contains(Math(4)))
+assert Math(1) in xs
+assert Math(4) not in xs
 
-    assert xs.remove(Math(1)) == MultiSet(Math(2), Math(3))
+egraph.check(xs.remove(Math(1)) == MultiSet(Math(2), Math(3)))
 
-    assert xs.length() == i64(3)
-    assert len(xs) == 3
+assert egraph.extract(xs.length()).eval() == 3
+assert len(xs) == 3
 
-    assert MultiSet(Math(1), Math(1)).length() == i64(2)
+egraph.check(MultiSet(Math(1), Math(1)).length() == i64(2))
 
-    assert MultiSet(Math(1)).pick() == Math(1)
+egraph.check(MultiSet(Math(1)).pick() == Math(1))
 
-    mapped = xs.map(square)
-    egraph.register(mapped)
-    egraph.run(math_ruleset)
-    assert mapped == MultiSet(Math(1), Math(4), Math(9))
+mapped = xs.map(square)
+egraph.register(mapped)
+egraph.run(math_ruleset)
+egraph.check(mapped == MultiSet(Math(1), Math(4), Math(9)))
 
-    assert xs + xs == MultiSet(Math(1), Math(2), Math(3), Math(1), Math(2), Math(3))
+egraph.check(xs + xs == MultiSet(Math(1), Math(2), Math(3), Math(1), Math(2), Math(3)))

egglog/exp/array_api.py

@@ -54,6 +54,7 @@ A: So that when converting from functional to lists we can use the same index fu
 
 from __future__ import annotations
 
+import contextlib
 import itertools
 import math
 import numbers
@@ -61,7 +62,7 @@ import os
 import sys
 from copy import copy
 from types import EllipsisType
-from typing import TYPE_CHECKING, ClassVar, TypeAlias, cast
+from typing import TYPE_CHECKING, Any, ClassVar, TypeAlias, cast
 
 import numpy as np
 
@@ -93,7 +94,7 @@ class Boolean(Expr, ruleset=array_api_ruleset):
 
     @method(preserve=True)
     def eval(self) -> bool:
-        return try_evaling(array_api_schedule, self, self.to_bool)
+        return try_evaling(_get_current_egraph(), array_api_schedule, self, self.to_bool)
 
     @property
     def to_bool(self) -> Bool: ...
@@ -218,7 +219,7 @@ class Int(Expr, ruleset=array_api_ruleset):
 
     @method(preserve=True)
     def eval(self) -> int:
-        return try_evaling(array_api_schedule, self, self.to_i64)
+        return try_evaling(_get_current_egraph(), array_api_schedule, self, self.to_i64)
 
     @method(preserve=True)
     def __index__(self) -> int:
@@ -330,7 +331,7 @@ class Float(Expr, ruleset=array_api_ruleset):
 
     @method(preserve=True)
     def eval(self) -> float:
-        return try_evaling(array_api_schedule, self, self.to_f64)
+        return try_evaling(_get_current_egraph(), array_api_schedule, self, self.to_f64)
 
     def abs(self) -> Float: ...
 
@@ -438,7 +439,7 @@ class TupleInt(Expr, ruleset=array_api_ruleset):
 
     @method(preserve=True)
     def eval(self) -> tuple[Int, ...]:
-        return try_evaling(array_api_schedule, self, self.to_vec)
+        return try_evaling(_get_current_egraph(), array_api_schedule, self, self.to_vec)
 
     def foldl(self, f: Callable[[Int, Int], Int], init: Int) -> Int: ...
     def foldl_boolean(self, f: Callable[[Boolean, Int], Boolean], init: Boolean) -> Boolean: ...
@@ -586,7 +587,7 @@ class TupleTupleInt(Expr, ruleset=array_api_ruleset):
 
     @method(preserve=True)
     def eval(self) -> tuple[TupleInt, ...]:
-        return try_evaling(array_api_schedule, self, self.to_vec)
+        return try_evaling(_get_current_egraph(), array_api_schedule, self, self.to_vec)
 
     def drop(self, n: Int) -> TupleTupleInt:
         return TupleTupleInt(self.length() - n, lambda i: self[i + n])
@@ -1282,7 +1283,7 @@ class TupleNDArray(Expr, ruleset=array_api_ruleset):
 
     @method(preserve=True)
     def eval(self) -> tuple[NDArray, ...]:
-        return try_evaling(array_api_schedule, self, self.to_vec)
+        return try_evaling(_get_current_egraph(), array_api_schedule, self, self.to_vec)
 
 
 converter(Vec[NDArray], TupleNDArray, lambda x: TupleNDArray.from_vec(x))
@@ -1941,3 +1942,42 @@ def array_api_vec_to_cons_ruleset(
 
 array_api_combined_ruleset = array_api_ruleset | array_api_vec_to_cons_ruleset
 array_api_schedule = array_api_combined_ruleset.saturate()
+
+_CURRENT_EGRAPH: None | EGraph = None
+
+
+@contextlib.contextmanager
+def set_array_api_egraph(egraph: EGraph) -> Iterator[None]:
+    """
+    Context manager that will set the current egraph. It will be set back after.
+    """
+    global _CURRENT_EGRAPH
+    assert _CURRENT_EGRAPH is None
+    _CURRENT_EGRAPH = egraph
+    yield
+    _CURRENT_EGRAPH = None
+
+
+def _get_current_egraph() -> EGraph:
+    return _CURRENT_EGRAPH or EGraph()
+
+
+def try_evaling(egraph: EGraph, schedule: Schedule, expr: Expr, prim_expr: BuiltinExpr) -> Any:
+    """
+    Try evaling the expression that will result in a primitive expression being fill.
+    if it fails, display the egraph and raise an error.
+    """
+    try:
+        extracted = egraph.extract(prim_expr)
+    except EggSmolError:
+        # If this primitive doesn't exist in the egraph, we need to try to create it by
+        # registering the expression and running the schedule
+        egraph.register(expr)
+        egraph.run(schedule)
+        try:
+            extracted = egraph.extract(prim_expr)
+        except BaseException as e:
+            # egraph.display(n_inline_leaves=1, split_primitive_outputs=True)
+            e.add_note(f"Cannot evaluate {egraph.extract(expr)}")
+            raise
+    return extracted.eval()  # type: ignore[attr-defined]

egglog/exp/array_api_jit.py

@@ -4,8 +4,8 @@ from typing import TypeVar, cast
 
 import numpy as np
 
-from egglog import EGraph, try_evaling
-from egglog.exp.array_api import NDArray
+from egglog import EGraph
+from egglog.exp.array_api import NDArray, set_array_api_egraph, try_evaling
 from egglog.exp.array_api_numba import array_api_numba_schedule
 from egglog.exp.array_api_program_gen import EvalProgram, array_api_program_gen_schedule, ndarray_function_two_program
 
@@ -20,8 +20,7 @@ def jit(fn: X) -> X:
     """
     egraph, res, res_optimized, program = function_to_program(fn, save_egglog_string=False)
     fn_program = EvalProgram(program, {"np": np})
-    with egraph.set_current():
-        fn = cast("X", try_evaling(array_api_program_gen_schedule, fn_program, fn_program.as_py_object))
+    fn = cast("X", try_evaling(egraph, array_api_program_gen_schedule, fn_program, fn_program.as_py_object))
     fn.initial_expr = res  # type: ignore[attr-defined]
     fn.expr = res_optimized  # type: ignore[attr-defined]
     return fn
@@ -32,9 +31,11 @@ def function_to_program(fn: Callable, save_egglog_string: bool) -> tuple[EGraph,
     arg1, arg2 = sig.parameters.keys()
     egraph = EGraph(save_egglog_string=save_egglog_string)
     with egraph:
-        with egraph.set_current():
+        with set_array_api_egraph(egraph):
             res = fn(NDArray.var(arg1), NDArray.var(arg2))
-        res_optimized = egraph.simplify(res, array_api_numba_schedule)
+        egraph.register(res)
+        egraph.run(array_api_numba_schedule)
+        res_optimized = egraph.extract(res)
 
     return (
         egraph,

{egglog-9.0.1.dist-info → egglog-10.0.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: egglog
-Version: 9.0.1
+Version: 10.0.0
 Classifier: Environment :: MacOS X
 Classifier: Environment :: Win32 (MS Windows)
 Classifier: Intended Audience :: Developers

{egglog-9.0.1.dist-info → egglog-10.0.0.dist-info}/RECORD

@@ -1,28 +1,28 @@
-egglog-9.0.1.dist-info/METADATA,sha256=c75HEwCDBdeEj8KEDeHe4FBweGGWS-mHSsYyOVRTDH4,4018
-egglog-9.0.1.dist-info/WHEEL,sha256=tAGdc4C2KTz7B2CZ8Jf3DcKSAviAbCg44UH9ma2gYww,96
-egglog-9.0.1.dist-info/licenses/LICENSE,sha256=TfaboMVZ81Q6OUaKjU7z6uVjSlcGKclLYcOpgDbm9_s,1091
+egglog-10.0.0.dist-info/METADATA,sha256=aWuZK6v1QbkwYpgfn9_9UnMiwCvPBONgvULMJCkMyFw,4019
+egglog-10.0.0.dist-info/WHEEL,sha256=tAGdc4C2KTz7B2CZ8Jf3DcKSAviAbCg44UH9ma2gYww,96
+egglog-10.0.0.dist-info/licenses/LICENSE,sha256=TfaboMVZ81Q6OUaKjU7z6uVjSlcGKclLYcOpgDbm9_s,1091
 egglog/bindings.pyi,sha256=fnRiGyHuM-m0fDe91kFryUC09U87f5XXPiY6Lno8bvk,14497
-egglog/builtins.py,sha256=_hgkpMUF-lCNasV7wG9r-7mSd-w0eNfyjYnwCKiWEP8,29143
+egglog/builtins.py,sha256=ZLXp5txd7CD_Wlb137DFdy_qGbvmvzbfb54tc_KbRfY,30692
 egglog/config.py,sha256=mALVaxh7zmGrbuyzaVKVmYKcu1lF703QsKJw8AF7gSM,176
-egglog/conversion.py,sha256=qhH-3hs9TYnQWooaJOiV_fnol19_Zzxxf6XZ7LHudM4,9895
+egglog/conversion.py,sha256=o6iE32taKz-PBX8aDo5i4jGELdEai-wrzMEZUC1lqsU,10257
 egglog/declarations.py,sha256=R4U8JhM6ChOMmPEgG6tL2sJ8wJGsIEEVVN5lVY-6k2w,23011
-egglog/egraph.py,sha256=t7Z4yaAeI1d-gQgsC3xDIE6U4OQQghN7AH3yqlqJ2gI,66259
+egglog/egraph.py,sha256=ssy2UHt8weA8GXzomqTWMKB389HzMo62k826gas5Fng,64537
 egglog/egraph_state.py,sha256=wkuHoXmEzcUFspvNl8ZG7_xS1J7XZTK3RPnlZEuvk9k,28720
-egglog/examples/bignum.py,sha256=pnXy8J8dQs-QjuLH4icahJFz2TyqRCL0gpBn9YPQg2M,543
+egglog/examples/bignum.py,sha256=r4UeVwpKK2iSZiT5MK4qqTWhWu76CAtSJ5YGDboO9tk,568
 egglog/examples/bool.py,sha256=pWZTjfXR1cFy3KcihLBU5AF5rn83ImORlhUUJ1YiAXc,733
-egglog/examples/eqsat_basic.py,sha256=ORXFYYEDsEZK2IPhHtoFsd-LdjMiQi1nn7kix4Nam0s,1011
+egglog/examples/eqsat_basic.py,sha256=ivYHFD52sjYyl5rWQ60ScxmvcbABYUBlPobDMu6HvN4,913
 egglog/examples/fib.py,sha256=wAn-PjazxgHDkXAU4o2xTk_GtM_iGL0biV66vWM1st4,520
 egglog/examples/higher_order_functions.py,sha256=5P-_10y8eTh-z3bze5GVzKHyFTeefTM774haF8KYGig,1164
 egglog/examples/lambda_.py,sha256=hQBOaSw_yorNcbkQVu2EhgSc0IZNWIny7asaOlcUk9s,8496
 egglog/examples/matrix.py,sha256=_zmjgfFr2O_LjTcsTD-45_38Y_M1sP3AV39K6oFxAdw,5136
-egglog/examples/multiset.py,sha256=4qVjs74SL_0T4MYOyr-3TomhMXuke2DVrQFk-CQJT04,1499
+egglog/examples/multiset.py,sha256=kSUFiN4qs8WXvexPCecqk3o5DzdywXacDmIiC5IDKCM,1507
 egglog/examples/ndarrays.py,sha256=T-wwef-n-3LDSjaO35zA8AZH5DXFFqq0XBSCQKEXV6E,4186
 egglog/examples/README.rst,sha256=QrbfmivODBvUvmY3-dHarcbC6bEvwoqAfTDhiI-aJxU,237
 egglog/examples/resolution.py,sha256=sKkbRI_v9XkQM0DriacKLINqKKDqYGFhvMCAS9tZbTA,2203
 egglog/examples/schedule_demo.py,sha256=iJtIbcLaZ7zK8UalY0z7KAKMqYjQx0MKTsNF24lKtik,652
 egglog/examples/__init__.py,sha256=KuhaJFOyz_rpUvEqZubsgLnv6rhQNE_AVFXA6bUnpdY,34
-egglog/exp/array_api.py,sha256=LDrNkrMznd937_It1sKn9H_gCnJMELC0jsGUjcjjA3Y,64166
-egglog/exp/array_api_jit.py,sha256=9eDVR0hCLeM3RLkIOrfitjX1bCvwoWZG1LF5nmRLWhc,1540
+egglog/exp/array_api.py,sha256=Y-mG8tC8ihFkeIbNDsDyZl7-7zP6YFxZGS4rfsWOG-8,65661
+egglog/exp/array_api_jit.py,sha256=c0lN-411M9LmEz8uM0ELXTtfUPMJDZFmy81-x3ck6PY,1591
 egglog/exp/array_api_loopnest.py,sha256=1KZ8kWtv00TVCbzqzWG92vz8Y1pBlPOLRIrJfxyDs9s,2518
 egglog/exp/array_api_numba.py,sha256=g0Yb7auB2MtJHgSzn7iVnsxn6ffL4FHggKvbEuuDwZs,2765
 egglog/exp/array_api_program_gen.py,sha256=qXaaW_Koz4P-KmQcyRHY1NDLCJL3O7NnE6iq2NX4yZU,22252
@@ -39,6 +39,6 @@ egglog/type_constraint_solver.py,sha256=kSJf8S1CxkicpPy_B2cNP9XiTf-d4TQNd7_IwfG-
 egglog/visualizer.css,sha256=DDmxZfLmFygwVkQwel0gGUiM5BZqgunawWeSfaRYBr8,259399
 egglog/visualizer.js,sha256=kDkYzE5cLLhmpNaWqSfQWzD0SlSsNe9KzT_GpNqhWgo,3789092
 egglog/visualizer_widget.py,sha256=TGkf61P_f4EvCFcyqwQC_8XYUkujFhcQnEG9vHWLPC4,1358
-egglog/__init__.py,sha256=YBsuQdA1JYMtdz79jVNGY1DXDPaWChjiba0Bu3RYfr0,276
-egglog/bindings.cp311-win_amd64.pyd,sha256=HGN2sZzhRRBh6b_yyDUaD1giNNn5F42QekvOd0kdpUA,4500480
-egglog-9.0.1.dist-info/RECORD,,
+egglog/__init__.py,sha256=TZ3IMVKoz5Qr97VFe-YT7oyT2NACikTis21PpVlSCmg,326
+egglog/bindings.cp311-win_amd64.pyd,sha256=lydYMT5DJB3jAsg_IpweX1BgjWbVm5Z7cPQDaakldrY,4500992
+egglog-10.0.0.dist-info/RECORD,,