egglog 9.0.1__cp311-cp311-manylinux_2_17_ppc64.manylinux2014_ppc64.whl → 10.0.1__cp311-cp311-manylinux_2_17_ppc64.manylinux2014_ppc64.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

@@ -5,6 +5,7 @@ Builtin sorts and function to egg.
 
 from __future__ import annotations
 
+from collections.abc import Callable
 from fractions import Fraction
 from functools import partial, reduce
 from types import FunctionType, MethodType
@@ -12,17 +13,15 @@ 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
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Iterator
+    from collections.abc import Iterator
 
 
 __all__ = [
@@ -32,6 +31,7 @@ __all__ = [
     "BigRatLike",
     "Bool",
     "BoolLike",
+    "BuiltinEvalError",
     "Map",
     "MapLike",
     "MultiSet",
@@ -56,6 +56,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 +83,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 +108,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 +143,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 +262,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 +351,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 +418,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 +482,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 +531,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 +618,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 +743,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 +852,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 +923,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 +1062,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
@@ -1839,12 +1817,13 @@ def _rewrite_or_rule_generator(gen: RewriteOrRuleGenerator, frame: FrameType) ->
     """
     Returns a thunk which will call the function with variables of the type and name of the arguments.
     """
-    # Get the local scope from where the function is defined, so that we can get any type hints that are in the scope
-    # but not in the globals
-    globals = gen.__globals__.copy()
-    if "Callable" not in globals:
-        globals["Callable"] = Callable
-    hints = get_type_hints(gen, globals, frame.f_locals)
+    # Need to manually pass in the frame locals from the generator, because otherwise classes defined within function
+    # will not be available in the annotations
+    # combine locals and globals so that they are the same dict. Otherwise get_type_hints will go through the wrong
+    # path and give an error for the test
+    # python/tests/test_no_import_star.py::test_no_import_star_rulesset
+    combined = {**gen.__globals__, **frame.f_locals}
+    hints = get_type_hints(gen, combined, combined)
     args = [_var(p.name, hints[p.name]) for p in signature(gen).parameters.values()]
     return list(gen(*args))  # type: ignore[misc]
 
@@ -1876,34 +1855,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/higher_order_functions.py

@@ -6,13 +6,10 @@ Higher Order Functions
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from collections.abc import Callable
 
 from egglog import *
 
-if TYPE_CHECKING:
-    from collections.abc import Callable
-
 
 class Math(Expr):
     def __init__(self, i: i64Like) -> None: ...

egglog/examples/lambda_.py

@@ -7,12 +7,11 @@ Lambda Calculus
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, ClassVar
+from collections.abc import Callable
+from typing import ClassVar
 
 from egglog import *
-
-if TYPE_CHECKING:
-    from collections.abc import Callable
+from egglog import Expr
 
 
 class Val(Expr):

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,14 +54,16 @@ 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
 import os
 import sys
+from collections.abc import Callable
 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
 
@@ -71,9 +73,10 @@ from egglog.runtime import RuntimeExpr
 from .program_gen import *
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Iterator
+    from collections.abc import Iterator
     from types import ModuleType
 
+
 # Pretend that exprs are numbers b/c sklearn does isinstance checks
 numbers.Integral.register(RuntimeExpr)
 
@@ -93,7 +96,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 +221,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 +333,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 +441,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 +589,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 +1285,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 +1944,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.1.dist-info}/METADATA

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

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

@@ -1,16 +1,16 @@
-egglog-9.0.1.dist-info/METADATA,sha256=avDfv1AMXRxm7eQcAhZYWbEm0UVuFB68FYUcT_jzUVQ,4008
-egglog-9.0.1.dist-info/WHEEL,sha256=TXoyImGYix7Rgz9vFJf5yHjHurpJZift2bEePFE703I,127
-egglog-9.0.1.dist-info/licenses/LICENSE,sha256=w7VlVv5O_FPZRo8Z-4Zb_q7D5ac3YDs8JUkMZ4Gq9CE,1070
+egglog-10.0.1.dist-info/METADATA,sha256=BKd5QxorELhavgKQf4E6fDVQ8FpCCifHy4v_BxFiXT0,4009
+egglog-10.0.1.dist-info/WHEEL,sha256=TXoyImGYix7Rgz9vFJf5yHjHurpJZift2bEePFE703I,127
+egglog-10.0.1.dist-info/licenses/LICENSE,sha256=w7VlVv5O_FPZRo8Z-4Zb_q7D5ac3YDs8JUkMZ4Gq9CE,1070
 egglog/visualizer_widget.py,sha256=LtVfzOtv2WeKtNuILQQ_9SOHWvRr8YdBYQDKQSgry_s,1319
 egglog/egraph_state.py,sha256=KxP4aXfh21vVqx1HbQMtb2PRxJg4Gzp1Po6Jp5Zppk0,28086
-egglog/builtins.py,sha256=yMvPxcYKjUxIFtDzKUr1XUwZSN5pLlcR66K7BCG5eS4,28095
+egglog/builtins.py,sha256=7cZk-MxUvOA3OT9RSX7WkYBoj1whAfp3tOB3a82ZTH0,29636
 egglog/exp/array_api_numba.py,sha256=X3H1TnCjPL92uVm6OvcWMJ11IeorAE58zWiOX6huPv4,2696
 egglog/exp/program_gen.py,sha256=9q8-ihkXLgeRU_BrINF3t3gThdS74KCzTo4UiRIR6Dk,13027
 egglog/exp/array_api_program_gen.py,sha256=0FpWZZJ3VqIod9vIYfQYUIiyswfj4h4K_4s6GZb8SJY,21742
 egglog/exp/siu_examples.py,sha256=yZ-sgH2Y12iTdwBUumP7D2OtCGL83M6pPW7PMobVFXc,719
 egglog/exp/__init__.py,sha256=nPtzrH1bz1LVZhZCuS0S9Qild8m5gEikjOVqWAFIa88,49
-egglog/exp/array_api.py,sha256=V-xiDokdp2QOOxduXwDO39BazxjG9mFMeD5tuIMhy-s,62223
-egglog/exp/array_api_jit.py,sha256=RzlMyEQ8pAKrQSjJ93jAHK-i3pFu6OXgs6ljDpPHqtQ,1496
+egglog/exp/array_api.py,sha256=jKQoDg3bYg0TGGEJLwXF3okWDkgi4JA4ayXzwf-9XB4,63706
+egglog/exp/array_api_jit.py,sha256=90RmyivRoCKzVtiYWnTBkK2q2FDkD2p1iq7vzbh68b4,1546
 egglog/exp/array_api_loopnest.py,sha256=-kbyorlGxvlaNsLx1nmLfEZHQM7VMEBwSKtV0l-bs0g,2444
 egglog/type_constraint_solver.py,sha256=_y52JJoxWJDzokvS07-QHJ4sQpvWE86_Dk43RAUqhAk,4516
 egglog/runtime.py,sha256=vyzN52njMucBaf8sRl4A5FVBrqDMfNaV8PlG8lzNxhc,25741
@@ -18,27 +18,27 @@ egglog/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 egglog/bindings.pyi,sha256=b0osaWlb6Yqh7Cw7I_DnRjU-6NTeXq1XOolvOvYMMcI,13830
 egglog/config.py,sha256=yM3FIcVCKnhWZmHD0pxkzx7ah7t5PxZx3WUqKtA9tjU,168
 egglog/declarations.py,sha256=UzpOjoZz21T9QfodmS3U00vx5VcTcryXmAw11CALix8,22193
-egglog/__init__.py,sha256=xVZPz0Pplz0OoXcrSPhNY7fHfja7Sao03DYbIKp8kcU,266
+egglog/__init__.py,sha256=5n2nas3VbxxgsvYeT8CutVNHbC2jggXQtDQ3RzX3vCM,315
 egglog/thunk.py,sha256=BwdqvNhY8g3KdSE5DtT1yIoxCF3mNGcqoODy8nXm9z8,2168
 egglog/pretty.py,sha256=5hcc5-tAGnhcpz_aClJ1PgE3StOvFqQRYXYFDWXbNP0,20476
-egglog/conversion.py,sha256=FnZFEZdZHkyrmg8aJ0_7SA-my5O0b2hjtjGGBrFR00U,9633
+egglog/conversion.py,sha256=lE9a3myOWyxlYbCVqOMKnqXw2_Ul24DJuJlVZ50HCA0,9989
 egglog/ipython_magic.py,sha256=2hs3g2cSiyDmbCvE2t1OINmu17Bb8MWV--2DpEWwO7I,1189
 egglog/functionalize.py,sha256=VMqiWcDbn1pYrY3FMxT143Bk90dMaxo2VapUL60-fvE,3847
-egglog/egraph.py,sha256=SJI2I7njTUdyg9BU0-SdwkHTvQVfck0mkCjA5zCvJ1c,64350
+egglog/egraph.py,sha256=Y2vAfJpH3ss51zc9krkfAVzkcttKtSwrnSpmrBupvXs,62867
 egglog/visualizer.css,sha256=eL0POoThQRc0P4OYnDT-d808ln9O5Qy6DizH9Z5LgWc,259398
 egglog/visualizer.js,sha256=2qZZ-9W_INJx4gZMYjnVXl27IjT_JNuQyEeI2dbjWoU,3753315
 egglog/examples/schedule_demo.py,sha256=JbXdPII7_adxtgyKVAiqCyV2sj88VZ-DhomYrdn8vuc,618
 egglog/examples/README.rst,sha256=ztTvpofR0eotSqGoCy_C1fPLDPCncjvcqDanXtLHNNU,232
 egglog/examples/fib.py,sha256=BOHxKWA7jGx4FURBmfmuZKfLo6xq9-uXAwAXjYid7LU,492
 egglog/examples/resolution.py,sha256=BJd5JClA3DBVGfiVRa-H0gbbFvIqeP3uYbhCXHblSQc,2119
-egglog/examples/lambda_.py,sha256=LK1ZSCITpvxpRMaqISZkRAKqNAU7QbaWq9oGe7ku2Ns,8208
-egglog/examples/multiset.py,sha256=bI7AmNlU3_D1GAhti2ojei5WQcbA3nnMun7y33hX0bE,1438
-egglog/examples/bignum.py,sha256=Rs_kttvEu_OJW26CRR0AY4Hup_FsvLpg6Klkt-5QfFU,512
+egglog/examples/lambda_.py,sha256=iQvwaXVhp2VNOMS7j1WwceZaiq3dqqilwUkMcW5GFBE,8194
+egglog/examples/multiset.py,sha256=fdEPvNFkHe_XmzBv90JCP8SCxoKgpg_CIDsR8HWcbnY,1447
+egglog/examples/bignum.py,sha256=LziQqBdOaYNhO3VQS6ZokHDO7_QSsOMV4Whyrj6-6nA,536
 egglog/examples/matrix.py,sha256=7_mPcMcgE-t_GJDyf76-nv3xhPIeN2mvFkc_p_Gnr8g,4961
 egglog/examples/__init__.py,sha256=wm9evUbMPfbtylXIjbDdRTAVMLH4OjT4Z77PCBFyaPU,31
-egglog/examples/higher_order_functions.py,sha256=hWs9sCM5mERuuYscw9Qp-W3HEHi0EMtAs57gV2Qv69k,1119
+egglog/examples/higher_order_functions.py,sha256=DNLIQfPJCX_DOLbHNiaYsfvcFIYCYOsRUqp99r9bpc8,1063
 egglog/examples/bool.py,sha256=e0z2YoYJsLlhpSADZK1yRYHzilyxSZWGiYAaM0DQ_Gw,695
 egglog/examples/ndarrays.py,sha256=mfr410eletH8gfdg-P8L90vlF6TUifvYV_-ryOwvZZE,4042
-egglog/examples/eqsat_basic.py,sha256=YKA_l6kKNfPk1geaf-qDohYHpvIji_HmABGr3Ez_N8k,966
-egglog/bindings.cpython-311-powerpc64-linux-gnu.so,sha256=GTNefRU7n9yaMMUINZsiHHm6WIiHNXX0Oagu_aDPhqw,6993280
-egglog-9.0.1.dist-info/RECORD,,
+egglog/examples/eqsat_basic.py,sha256=2xtM81gG9Br72mr58N-2BUeksR7C_UXnZJ4MvzSPplc,869
+egglog/bindings.cpython-311-powerpc64-linux-gnu.so,sha256=El3R-hSWj7dW_ACcndW3WdEwQ6PeTC_uHNxnf9MVVTA,6992704
+egglog-10.0.1.dist-info/RECORD,,