egglog 10.0.2__cp310-cp310-win_amd64.whl → 11.0.0__cp310-cp310-win_amd64.whl

egglog/egraph.py

@@ -5,7 +5,7 @@ import inspect
 import pathlib
 import tempfile
 from collections.abc import Callable, Generator, Iterable
-from contextvars import ContextVar
+from contextvars import ContextVar, Token
 from dataclasses import InitVar, dataclass, field
 from functools import partial
 from inspect import Parameter, currentframe, signature
@@ -29,6 +29,7 @@ from typing_extensions import Never, ParamSpec, Self, Unpack, assert_never
 
 from . import bindings
 from .conversion import *
+from .conversion import convert_to_same_type, resolve_literal
 from .declarations import *
 from .egraph_state import *
 from .ipython_magic import IN_IPYTHON
@@ -82,7 +83,6 @@ __all__ = [
     "run",
     "seq",
     "set_",
-    "simplify",
     "subsume",
     "union",
     "unstable_combine_rulesets",
@@ -110,12 +110,12 @@ IGNORED_ATTRIBUTES = {
     "__weakref__",
     "__orig_bases__",
     "__annotations__",
-    "__hash__",
     "__qualname__",
     "__firstlineno__",
     "__static_attributes__",
+    "__match_args__",
     # Ignore all reflected binary method
-    *REFLECTED_BINARY_METHODS.keys(),
+    *(f"__r{m[2:]}" for m in NUMERIC_BINARY_METHODS),
 }
 
 
@@ -139,15 +139,6 @@ ALWAYS_PRESERVED = {
 }
 
 
-def simplify(x: EXPR, schedule: Schedule | None = None) -> EXPR:
-    """
-    Simplify an expression by running the schedule.
-    """
-    if schedule:
-        return EGraph().simplify(x, schedule)
-    return EGraph().extract(x)
-
-
 def check_eq(x: BASE_EXPR, y: BASE_EXPR, schedule: Schedule | None = None, *, add_second=True, display=False) -> EGraph:
     """
     Verifies that two expressions are equal after running the schedule.
@@ -291,7 +282,6 @@ def function(
     mutates_first_arg: bool = ...,
     unextractable: bool = ...,
     ruleset: Ruleset | None = ...,
-    use_body_as_name: bool = ...,
     subsume: bool = ...,
 ) -> Callable[[CONSTRUCTOR_CALLABLE], CONSTRUCTOR_CALLABLE]: ...
 
@@ -371,6 +361,7 @@ class BaseExpr(metaclass=_ExprMetaclass):
 
     def __ne__(self, other: Self) -> Unit: ...  # type: ignore[override, empty-body]
 
+    # not currently dissalowing other types of equality https://github.com/python/typeshed/issues/8217#issuecomment-3140873292
     def __eq__(self, other: Self) -> Fact: ...  # type: ignore[override, empty-body]
 
 
@@ -404,7 +395,7 @@ def _generate_class_decls(  # noqa: C901,PLR0912
     )
     type_vars = tuple(ClassTypeVarRef.from_type_var(p) for p in parameters)
     del parameters
-    cls_decl = ClassDecl(egg_sort, type_vars, builtin)
+    cls_decl = ClassDecl(egg_sort, type_vars, builtin, match_args=namespace.pop("__match_args__", ()))
     decls = Declarations(_classes={cls_name: cls_decl})
     # Update class think eagerly when resolving so that lookups work in methods
     runtime_cls.__egg_decls_thunk__ = Thunk.value(decls)
@@ -456,6 +447,9 @@ def _generate_class_decls(  # noqa: C901,PLR0912
             continue
         locals = frame.f_locals
         ref: ClassMethodRef | MethodRef | PropertyRef | InitRef
+        # TODO: Store deprecated message so we can print at runtime
+        if (getattr(fn, "__deprecated__", None)) is not None:
+            fn = fn.__wrapped__  # type: ignore[attr-defined]
         match fn:
             case classmethod():
                 ref = ClassMethodRef(cls_name, method_name)
@@ -477,7 +471,7 @@ def _generate_class_decls(  # noqa: C901,PLR0912
             decls.set_function_decl(ref, decl)
             continue
         try:
-            _, add_rewrite = _fn_decl(
+            add_rewrite = _fn_decl(
                 decls,
                 egg_fn,
                 ref,
@@ -515,7 +509,6 @@ class _FunctionConstructor:
     merge: Callable[[object, object], object] | None = None
     unextractable: bool = False
     ruleset: Ruleset | None = None
-    use_body_as_name: bool = False
     subsume: bool = False
 
     def __call__(self, fn: Callable) -> RuntimeFunction:
@@ -523,11 +516,10 @@ class _FunctionConstructor:
 
     def create_decls(self, fn: Callable) -> tuple[Declarations, CallableRef]:
         decls = Declarations()
-        ref = None if self.use_body_as_name else FunctionRef(fn.__name__)
-        ref, add_rewrite = _fn_decl(
+        add_rewrite = _fn_decl(
             decls,
             self.egg_fn,
-            ref,
+            ref := FunctionRef(fn.__name__),
             fn,
             self.hint_locals,
             self.cost,
@@ -545,8 +537,7 @@ class _FunctionConstructor:
 def _fn_decl(
     decls: Declarations,
     egg_name: str | None,
-    # If ref is Callable, then generate the ref from the function name
-    ref: FunctionRef | MethodRef | PropertyRef | ClassMethodRef | InitRef | None,
+    ref: FunctionRef | MethodRef | PropertyRef | ClassMethodRef | InitRef,
     fn: object,
     # Pass in the locals, retrieved from the frame when wrapping,
     # so that we support classes and function defined inside of other functions (which won't show up in the globals)
@@ -559,7 +550,7 @@ def _fn_decl(
     ruleset: Ruleset | None = None,
     unextractable: bool = False,
     reverse_args: bool = False,
-) -> tuple[CallableRef, Callable[[], None]]:
+) -> Callable[[], None]:
     """
     Sets the function decl for the function object and returns the ref as well as a thunk that sets the default callable.
     """
@@ -619,8 +610,8 @@ def _fn_decl(
         else resolve_literal(
             return_type,
             merge(
-                RuntimeExpr.__from_values__(decls, TypedExprDecl(return_type.to_just(), VarDecl("old", False))),
-                RuntimeExpr.__from_values__(decls, TypedExprDecl(return_type.to_just(), VarDecl("new", False))),
+                RuntimeExpr.__from_values__(decls, TypedExprDecl(return_type.to_just(), UnboundVarDecl("old", "old"))),
+                RuntimeExpr.__from_values__(decls, TypedExprDecl(return_type.to_just(), UnboundVarDecl("new", "new"))),
             ),
             lambda: decls,
         )
@@ -628,51 +619,40 @@ def _fn_decl(
     decls |= merged
 
     # defer this in generator so it doesn't resolve for builtins eagerly
-    args = (TypedExprDecl(tp.to_just(), VarDecl(name, False)) for name, tp in zip(arg_names, arg_types, strict=True))
-    res_ref: FunctionRef | MethodRef | ClassMethodRef | PropertyRef | InitRef | UnnamedFunctionRef
-    res_thunk: Callable[[], object]
-    # If we were not passed in a ref, this is an unnamed funciton, so eagerly compute the value and use that to refer to it
-    if not ref:
-        tuple_args = tuple(args)
-        res = _create_default_value(decls, ref, fn, tuple_args, ruleset)
-        assert isinstance(res, RuntimeExpr)
-        res_ref = UnnamedFunctionRef(tuple_args, res.__egg_typed_expr__)
-        decls._unnamed_functions.add(res_ref)
-        res_thunk = Thunk.value(res)
+    args = (TypedExprDecl(tp.to_just(), UnboundVarDecl(name)) for name, tp in zip(arg_names, arg_types, strict=True))
 
+    return_type_is_eqsort = (
+        not decls._classes[return_type.name].builtin if isinstance(return_type, TypeRefWithVars) else False
+    )
+    is_constructor = not is_builtin and return_type_is_eqsort and merged is None
+    signature_ = FunctionSignature(
+        return_type=None if mutates_first_arg else return_type,
+        var_arg_type=var_arg_type,
+        arg_types=arg_types,
+        arg_names=arg_names,
+        arg_defaults=tuple(a.__egg_typed_expr__.expr if a is not None else None for a in arg_defaults),
+        reverse_args=reverse_args,
+    )
+    decl: ConstructorDecl | FunctionDecl
+    if is_constructor:
+        decl = ConstructorDecl(signature_, egg_name, cost, unextractable)
     else:
-        return_type_is_eqsort = (
-            not decls._classes[return_type.name].builtin if isinstance(return_type, TypeRefWithVars) else False
-        )
-        is_constructor = not is_builtin and return_type_is_eqsort and merged is None
-        signature_ = FunctionSignature(
-            return_type=None if mutates_first_arg else return_type,
-            var_arg_type=var_arg_type,
-            arg_types=arg_types,
-            arg_names=arg_names,
-            arg_defaults=tuple(a.__egg_typed_expr__.expr if a is not None else None for a in arg_defaults),
-            reverse_args=reverse_args,
+        if cost is not None:
+            msg = "Cost can only be set for constructors"
+            raise ValueError(msg)
+        if unextractable:
+            msg = "Unextractable can only be set for constructors"
+            raise ValueError(msg)
+        decl = FunctionDecl(
+            signature=signature_,
+            egg_name=egg_name,
+            merge=merged.__egg_typed_expr__.expr if merged is not None else None,
+            builtin=is_builtin,
         )
-        decl: ConstructorDecl | FunctionDecl
-        if is_constructor:
-            decl = ConstructorDecl(signature_, egg_name, cost, unextractable)
-        else:
-            if cost is not None:
-                msg = "Cost can only be set for constructors"
-                raise ValueError(msg)
-            if unextractable:
-                msg = "Unextractable can only be set for constructors"
-                raise ValueError(msg)
-            decl = FunctionDecl(
-                signature=signature_,
-                egg_name=egg_name,
-                merge=merged.__egg_typed_expr__.expr if merged is not None else None,
-                builtin=is_builtin,
-            )
-        res_ref = ref
-        decls.set_function_decl(ref, decl)
-        res_thunk = Thunk.fn(_create_default_value, decls, ref, fn, args, ruleset, context=f"creating {ref}")
-    return res_ref, Thunk.fn(_add_default_rewrite_function, decls, res_ref, return_type, ruleset, res_thunk, subsume)
+    decls.set_function_decl(ref, decl)
+    return Thunk.fn(
+        _add_default_rewrite_function, decls, ref, fn, args, ruleset, subsume, return_type, context=f"creating {ref}"
+    )
 
 
 # Overload to support aritys 0-4 until variadic generic support map, so we can map from type to value
@@ -707,7 +687,7 @@ def relation(name: str, /, *tps: type, egg_fn: str | None = None) -> Callable[..
 
 
 def _relation_decls(name: str, tps: tuple[type, ...], egg_fn: str | None) -> Declarations:
-    from .builtins import Unit
+    from .builtins import Unit  # noqa: PLC0415
 
     decls = Declarations()
     decls |= cast("RuntimeClass", Unit)
@@ -746,13 +726,15 @@ def _constant_thunk(
     return decls, TypedExprDecl(type_ref.to_just(), CallDecl(callable_ref))
 
 
-def _create_default_value(
+def _add_default_rewrite_function(
     decls: Declarations,
-    ref: CallableRef | None,
+    ref: FunctionRef | MethodRef | PropertyRef | ClassMethodRef | InitRef,
     fn: Callable,
     args: Iterable[TypedExprDecl],
     ruleset: Ruleset | None,
-) -> object:
+    subsume: bool,
+    res_type: TypeOrVarRef,
+) -> None:
     args: list[object] = [RuntimeExpr.__from_values__(decls, a) for a in args]
 
     # If this is a classmethod, add the class as the first arg
@@ -760,21 +742,8 @@ def _create_default_value(
         tp = decls.get_paramaterized_class(ref.class_name)
         args.insert(0, RuntimeClass(Thunk.value(decls), tp))
     with set_current_ruleset(ruleset):
-        return fn(*args)
-
-
-def _add_default_rewrite_function(
-    decls: Declarations,
-    ref: CallableRef,
-    res_type: TypeOrVarRef,
-    ruleset: Ruleset | None,
-    value_thunk: Callable[[], object],
-    subsume: bool,
-) -> None:
-    """
-    Helper functions that resolves a value thunk to create the default value.
-    """
-    _add_default_rewrite(decls, ref, res_type, value_thunk(), ruleset, subsume)
+        res = fn(*args)
+    _add_default_rewrite(decls, ref, res_type, res, ruleset, subsume)
 
 
 def _add_default_rewrite(
@@ -794,6 +763,13 @@ def _add_default_rewrite(
         return
     resolved_value = resolve_literal(type_ref, default_rewrite, Thunk.value(decls))
     rewrite_decl = DefaultRewriteDecl(ref, resolved_value.__egg_typed_expr__.expr, subsume)
+    ruleset_decls = _add_default_rewrite_inner(decls, rewrite_decl, ruleset)
+    ruleset_decls |= resolved_value
+
+
+def _add_default_rewrite_inner(
+    decls: Declarations, rewrite_decl: DefaultRewriteDecl, ruleset: Ruleset | None
+) -> Declarations:
     if ruleset:
         ruleset_decls = ruleset._current_egg_decls
         ruleset_decl = ruleset.__egg_ruleset__
@@ -801,7 +777,7 @@ def _add_default_rewrite(
         ruleset_decls = decls
         ruleset_decl = decls.default_ruleset
     ruleset_decl.rules.append(rewrite_decl)
-    ruleset_decls |= resolved_value
+    return ruleset_decls
 
 
 def _last_param_variable(params: list[Parameter]) -> bool:
@@ -882,6 +858,7 @@ class EGraph:
         self._add_decls(decls)
         return self._state.callable_ref_to_egg(ref)[0]
 
+    # TODO: Change let to be action...
     def let(self, name: str, expr: BASE_EXPR) -> BASE_EXPR:
         """
         Define a new expression in the egraph and return a reference to it.
@@ -893,38 +870,10 @@ class EGraph:
         return cast(
             "BASE_EXPR",
             RuntimeExpr.__from_values__(
-                self.__egg_decls__, TypedExprDecl(runtime_expr.__egg_typed_expr__.tp, VarDecl(name, True))
+                self.__egg_decls__, TypedExprDecl(runtime_expr.__egg_typed_expr__.tp, LetRefDecl(name))
             ),
         )
 
-    @overload
-    def simplify(self, expr: BASE_EXPR, limit: int, /, *until: Fact, ruleset: Ruleset | None = None) -> BASE_EXPR: ...
-
-    @overload
-    def simplify(self, expr: BASE_EXPR, schedule: Schedule, /) -> BASE_EXPR: ...
-
-    def simplify(
-        self, expr: BASE_EXPR, limit_or_schedule: int | Schedule, /, *until: Fact, ruleset: Ruleset | None = None
-    ) -> BASE_EXPR:
-        """
-        Simplifies the given expression.
-        """
-        schedule = run(ruleset, *until) * limit_or_schedule if isinstance(limit_or_schedule, int) else limit_or_schedule
-        del limit_or_schedule, until, ruleset
-        runtime_expr = to_runtime_expr(expr)
-        self._add_decls(runtime_expr, schedule)
-        egg_schedule = self._state.schedule_to_egg(schedule.schedule)
-        typed_expr = runtime_expr.__egg_typed_expr__
-        # Must also register type
-        egg_expr = self._state.typed_expr_to_egg(typed_expr)
-        self._egraph.run_program(bindings.Simplify(span(1), egg_expr, egg_schedule))
-        extract_report = self._egraph.extract_report()
-        if not isinstance(extract_report, bindings.Best):
-            msg = "No extract report saved"
-            raise ValueError(msg)  # noqa: TRY004
-        (new_typed_expr,) = self._state.exprs_from_egg(extract_report.termdag, [extract_report.term], typed_expr.tp)
-        return cast("BASE_EXPR", RuntimeExpr.__from_values__(self.__egg_decls__, new_typed_expr))
-
     def include(self, path: str) -> None:
         """
         Include a file of rules.
@@ -1036,9 +985,7 @@ class EGraph:
         self._add_decls(expr)
         expr = self._state.typed_expr_to_egg(expr.__egg_typed_expr__)
         try:
-            self._egraph.run_program(
-                bindings.ActionCommand(bindings.Extract(span(2), expr, bindings.Lit(span(2), bindings.Int(n))))
-            )
+            self._egraph.run_program(bindings.Extract(span(2), expr, bindings.Lit(span(2), bindings.Int(n))))
         except BaseException as e:
             raise add_note("Extracting: " + str(expr), e)  # noqa: B904
         extract_report = self._egraph.extract_report()
@@ -1138,9 +1085,9 @@ class EGraph:
 
         If in IPython it will display it inline, otherwise it will write it to a file and open it.
         """
-        from IPython.display import SVG, display
+        from IPython.display import SVG, display  # noqa: PLC0415
 
-        from .visualizer_widget import VisualizerWidget
+        from .visualizer_widget import VisualizerWidget  # noqa: PLC0415
 
         if graphviz:
             if IN_IPYTHON:
@@ -1167,7 +1114,7 @@ class EGraph:
 
         If an `expr` is passed, it's also extracted after each run and printed
         """
-        from .visualizer_widget import VisualizerWidget
+        from .visualizer_widget import VisualizerWidget  # noqa: PLC0415
 
         def to_json() -> str:
             if expr is not None:
@@ -1559,16 +1506,17 @@ def rule(*facts: FactLike, ruleset: None = None, name: str | None = None) -> _Ru
     return _RuleBuilder(facts=_fact_likes(facts), name=name, ruleset=ruleset)
 
 
-def var(name: str, bound: type[T]) -> T:
+def var(name: str, bound: type[T], egg_name: str | None = None) -> T:
     """Create a new variable with the given name and type."""
-    return cast("T", _var(name, bound))
+    return cast("T", _var(name, bound, egg_name=egg_name))
 
 
-def _var(name: str, bound: object) -> RuntimeExpr:
+def _var(name: str, bound: object, egg_name: str | None) -> RuntimeExpr:
     """Create a new variable with the given name and type."""
     decls_like, type_ref = resolve_type_annotation(bound)
     return RuntimeExpr(
-        Thunk.fn(Declarations.create, decls_like), Thunk.value(TypedExprDecl(type_ref.to_just(), VarDecl(name, False)))
+        Thunk.fn(Declarations.create, decls_like),
+        Thunk.value(TypedExprDecl(type_ref.to_just(), UnboundVarDecl(name, egg_name))),
     )
 
 
@@ -1659,7 +1607,7 @@ class _NeBuilder(Generic[BASE_EXPR]):
     lhs: BASE_EXPR
 
     def to(self, rhs: BASE_EXPR) -> Unit:
-        from .builtins import Unit
+        from .builtins import Unit  # noqa: PLC0415
 
         lhs = to_runtime_expr(self.lhs)
         rhs = convert_to_same_type(rhs, lhs)
@@ -1818,7 +1766,7 @@ def _rewrite_or_rule_generator(gen: RewriteOrRuleGenerator, frame: FrameType) ->
     # 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()]
+    args = [_var(p.name, hints[p.name], egg_name=None) for p in signature(gen).parameters.values()]
     return list(gen(*args))  # type: ignore[misc]
 
 
@@ -1844,7 +1792,7 @@ def get_current_ruleset() -> Ruleset | None:
 
 @contextlib.contextmanager
 def set_current_ruleset(r: Ruleset | None) -> Generator[None, None, None]:
-    token = _CURRENT_RULESET.set(r)
+    token: Token[Ruleset | None] = _CURRENT_RULESET.set(r)
     try:
         yield
     finally:

egglog/egraph_state.py

@@ -108,7 +108,7 @@ class EGraphState:
             case RulesetDecl(rules):
                 if name not in self.rulesets:
                     if name:
-                        self.egraph.run_program(bindings.AddRuleset(name))
+                        self.egraph.run_program(bindings.AddRuleset(span(), name))
                     added_rules = self.rulesets[name] = set()
                 else:
                     added_rules = self.rulesets[name]
@@ -125,7 +125,7 @@ class EGraphState:
                 self.rulesets[name] = set()
                 for ruleset in rulesets:
                     self.ruleset_to_egg(ruleset)
-                self.egraph.run_program(bindings.UnstableCombinedRuleset(name, list(rulesets)))
+                self.egraph.run_program(bindings.UnstableCombinedRuleset(span(), name, list(rulesets)))
 
     def command_to_egg(self, cmd: CommandDecl, ruleset: str) -> bindings._Command | None:
         match cmd:
@@ -160,7 +160,7 @@ class EGraphState:
                 assert isinstance(sig, FunctionSignature)
                 # Replace args with rule_var_name mapping
                 arg_mapping = tuple(
-                    TypedExprDecl(tp.to_just(), VarDecl(name, False))
+                    TypedExprDecl(tp.to_just(), UnboundVarDecl(name))
                     for name, tp in zip(sig.arg_names, sig.arg_types, strict=False)
                 )
                 rewrite_decl = RewriteDecl(
@@ -179,7 +179,7 @@ class EGraphState:
     def action_to_egg(self, action: ActionDecl, expr_to_let: bool = False) -> bindings._Action | None:  # noqa: C901, PLR0911, PLR0912
         match action:
             case LetDecl(name, typed_expr):
-                var_decl = VarDecl(name, True)
+                var_decl = LetRefDecl(name)
                 var_egg = self._expr_to_egg(var_decl)
                 self.expr_to_egg_cache[var_decl] = var_egg
                 return bindings.Let(span(), var_egg.name, self.typed_expr_to_egg(typed_expr))
@@ -369,7 +369,8 @@ class EGraphState:
         """
         Rewrites this expression as a let binding if it's not already a let binding.
         """
-        var_decl = VarDecl(f"__expr_{hash(typed_expr)}", True)
+        # TODO: Replace with counter so that it works with hash collisions and is more stable
+        var_decl = LetRefDecl(f"__expr_{hash(typed_expr)}")
         if var_decl in self.expr_to_egg_cache:
             return None
         var_egg = self._expr_to_egg(var_decl)
@@ -387,7 +388,7 @@ class EGraphState:
     def _expr_to_egg(self, expr_decl: CallDecl) -> bindings.Call: ...
 
     @overload
-    def _expr_to_egg(self, expr_decl: VarDecl) -> bindings.Var: ...
+    def _expr_to_egg(self, expr_decl: UnboundVarDecl | LetRefDecl) -> bindings.Var: ...
 
     @overload
     def _expr_to_egg(self, expr_decl: ExprDecl) -> bindings._Expr: ...
@@ -402,11 +403,10 @@ class EGraphState:
             pass
         res: bindings._Expr
         match expr_decl:
-            case VarDecl(name, is_let):
-                # prefix let bindings with % to avoid name conflicts with rewrites
-                if is_let:
-                    name = f"%{name}"
-                res = bindings.Var(span(), name)
+            case LetRefDecl(name):
+                res = bindings.Var(span(), f"{name}")
+            case UnboundVarDecl(name, egg_name):
+                res = bindings.Var(span(), egg_name or f"_{name}")
             case LitDecl(value):
                 l: bindings._Literal
                 match value:
@@ -467,7 +467,8 @@ class EGraphState:
                 return name
 
             case ConstantRef(name):
-                return name
+                # Prefix to avoid name collisions with local vars
+                return f"%{name}"
             case (
                 MethodRef(cls_name, name)
                 | ClassMethodRef(cls_name, name)
@@ -549,7 +550,7 @@ class FromEggState:
         """
         expr_decl: ExprDecl
         if isinstance(term, bindings.TermVar):
-            expr_decl = VarDecl(term.name, True)
+            expr_decl = LetRefDecl(term.name)
         elif isinstance(term, bindings.TermLit):
             value = term.value
             expr_decl = LitDecl(None if isinstance(value, bindings.Unit) else value.value)
@@ -624,7 +625,9 @@ class FromEggState:
                     # but dont need to store them
                     bound_tp_params if isinstance(callable_ref, ClassMethodRef | InitRef) else None,
                 )
-        raise ValueError(f"Could not find callable ref for call {term}")
+        raise ValueError(
+            f"Could not find callable ref for call {term}. None of these refs matched the types: {self.state.egg_fn_to_callable_refs[term.name]}"
+        )
 
     def resolve_term(self, term_id: int, tp: JustTypeRef) -> TypedExprDecl:
         try:

egglog/examples/bignum.py

@@ -14,7 +14,7 @@ z = BigRat(x, y)
 
 egraph = EGraph()
 
-assert egraph.extract(z.numer.to_string()).eval() == "-617"
+assert egraph.extract(z.numer.to_string()).value == "-617"
 
 
 @function

egglog/examples/multiset.py

@@ -32,7 +32,7 @@ egraph.register(xs)
 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(egraph.extract(xs).eval()) == Counter({Math(1): 1, Math(2): 1, Math(3): 1})
+assert Counter(egraph.extract(xs).value) == Counter({Math(1): 1, Math(2): 1, Math(3): 1})
 
 
 inserted = MultiSet(Math(1), Math(2), Math(3), Math(4))
@@ -45,7 +45,7 @@ assert Math(4) not in xs
 
 egraph.check(xs.remove(Math(1)) == MultiSet(Math(2), Math(3)))
 
-assert egraph.extract(xs.length()).eval() == 3
+assert egraph.extract(xs.length()).value == 3
 assert len(xs) == 3
 
 egraph.check(MultiSet(Math(1), Math(1)).length() == i64(2))

egglog/exp/array_api.py

@@ -154,6 +154,18 @@ class Int(Expr, ruleset=array_api_ruleset):
     def __eq__(self, other: IntLike) -> Boolean:  # type: ignore[override]
         ...
 
+    # add a hash so that this test can pass
+    # https://github.com/scikit-learn/scikit-learn/blob/6fd23fca53845b32b249f2b36051c081b65e2fab/sklearn/utils/validation.py#L486-L487
+    @method(preserve=True)
+    def __hash__(self) -> int:
+        egraph = _get_current_egraph()
+        egraph.register(self)
+        egraph.run(array_api_schedule)
+        simplified = egraph.extract(self)
+        return hash(cast("RuntimeExpr", simplified).__egg_typed_expr__)
+
+    def __round__(self, ndigits: OptionalIntLike = None) -> Int: ...
+
     # TODO: Fix this?
     # Make != always return a Bool, so that numpy.unique works on a tuple of ints
     # In _unique1d
@@ -280,6 +292,8 @@ def _int(i: i64, j: i64, r: Boolean, o: Int, b: Int):
     yield rewrite(Int.if_(TRUE, o, b), subsume=True).to(o)
     yield rewrite(Int.if_(FALSE, o, b), subsume=True).to(b)
 
+    yield rewrite(o.__round__(OptionalInt.none)).to(o)
+
     # Never cannot be equal to anything real
     yield rule(eq(Int.NEVER).to(Int(i))).then(panic("Int.NEVER cannot be equal to any real int"))
 
@@ -354,8 +368,14 @@ class Float(Expr, ruleset=array_api_ruleset):
     def __sub__(self, other: FloatLike) -> Float: ...
 
     def __pow__(self, other: FloatLike) -> Float: ...
+    def __round__(self, ndigits: OptionalIntLike = None) -> Float: ...
 
     def __eq__(self, other: FloatLike) -> Boolean: ...  # type: ignore[override]
+    def __ne__(self, other: FloatLike) -> Boolean: ...  # type: ignore[override]
+    def __lt__(self, other: FloatLike) -> Boolean: ...
+    def __le__(self, other: FloatLike) -> Boolean: ...
+    def __gt__(self, other: FloatLike) -> Boolean: ...
+    def __ge__(self, other: FloatLike) -> Boolean: ...
 
 
 converter(float, Float, lambda x: Float(x))
@@ -366,9 +386,10 @@ FloatLike: TypeAlias = Float | float | IntLike
 
 
 @array_api_ruleset.register
-def _float(fl: Float, f: f64, f2: f64, i: i64, r: BigRat, r1: BigRat):
+def _float(fl: Float, f: f64, f2: f64, i: i64, r: BigRat, r1: BigRat, i_: Int):
     return [
         rule(eq(fl).to(Float(f))).then(set_(fl.to_f64).to(f)),
+        rewrite(Float.from_int(Int(i))).to(Float(f64.from_i64(i))),
         rewrite(Float(f).abs()).to(Float(f), f >= 0.0),
         rewrite(Float(f).abs()).to(Float(-f), f < 0.0),
         # Convert from float to rationl, if its a whole number i.e. can be converted to int
@@ -383,11 +404,22 @@ def _float(fl: Float, f: f64, f2: f64, i: i64, r: BigRat, r1: BigRat):
         rewrite(Float.rational(r) - Float.rational(r1)).to(Float.rational(r - r1)),
         rewrite(Float.rational(r) * Float.rational(r1)).to(Float.rational(r * r1)),
         rewrite(Float(f) ** Float(f2)).to(Float(f**f2)),
-        # ==
+        # comparisons
         rewrite(Float(f) == Float(f)).to(TRUE),
         rewrite(Float(f) == Float(f2)).to(FALSE, ne(f).to(f2)),
+        rewrite(Float(f) != Float(f2)).to(TRUE, f != f2),
+        rewrite(Float(f) != Float(f)).to(FALSE),
+        rewrite(Float(f) >= Float(f2)).to(TRUE, f >= f2),
+        rewrite(Float(f) >= Float(f2)).to(FALSE, f < f2),
+        rewrite(Float(f) <= Float(f2)).to(TRUE, f <= f2),
+        rewrite(Float(f) <= Float(f2)).to(FALSE, f > f2),
+        rewrite(Float(f) > Float(f2)).to(TRUE, f > f2),
+        rewrite(Float(f) > Float(f2)).to(FALSE, f <= f2),
+        rewrite(Float(f) < Float(f2)).to(TRUE, f < f2),
         rewrite(Float.rational(r) == Float.rational(r)).to(TRUE),
         rewrite(Float.rational(r) == Float.rational(r1)).to(FALSE, ne(r).to(r1)),
+        # round
+        rewrite(Float.rational(r).__round__()).to(Float.rational(r.round())),
     ]
 
 
@@ -671,6 +703,8 @@ class OptionalInt(Expr, ruleset=array_api_ruleset):
     def some(cls, value: Int) -> OptionalInt: ...
 
 
+OptionalIntLike: TypeAlias = OptionalInt | IntLike | None
+
 converter(type(None), OptionalInt, lambda _: OptionalInt.none)
 converter(Int, OptionalInt, OptionalInt.some)
 
@@ -1982,4 +2016,4 @@ def try_evaling(egraph: EGraph, schedule: Schedule, expr: Expr, prim_expr: Built
         except BaseException as e:
             # egraph.display(n_inline_leaves=1, split_primitive_outputs=True)
             raise add_note(f"Cannot evaluate {egraph.extract(expr)}", e)  # noqa: B904
-    return extracted.eval()  # type: ignore[attr-defined]
+    return extracted.value  # type: ignore[attr-defined]

egglog/exp/array_api_jit.py

@@ -14,16 +14,22 @@ from .program_gen import Program
 X = TypeVar("X", bound=Callable)
 
 
-def jit(fn: X) -> X:
+def jit(
+    fn: X,
+    *,
+    handle_expr: Callable[[NDArray], None] | None = None,
+    handle_optimized_expr: Callable[[NDArray], None] | None = None,
+) -> X:
     """
     Jit compiles a function
     """
     egraph, res, res_optimized, program = function_to_program(fn, save_egglog_string=False)
+    if handle_expr:
+        handle_expr(res)
+    if handle_optimized_expr:
+        handle_optimized_expr(res_optimized)
     fn_program = EvalProgram(program, {"np": np})
-    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
+    return cast("X", try_evaling(egraph, array_api_program_gen_schedule, fn_program, fn_program.as_py_object))
 
 
 def function_to_program(fn: Callable, save_egglog_string: bool) -> tuple[EGraph, NDArray, NDArray, Program]: