egglog 7.1.0__cp310-none-win_amd64.whl → 7.2.0__cp310-none-win_amd64.whl

egglog/declarations.py

@@ -71,6 +71,8 @@ __all__ = [
     "CommandDecl",
     "SpecialFunctions",
     "FunctionSignature",
+    "DefaultRewriteDecl",
+    "InitRef",
 ]
 
 
@@ -80,7 +82,13 @@ class DelayedDeclerations:
 
     @property
     def __egg_decls__(self) -> Declarations:
-        return self.__egg_decls_thunk__()
+        try:
+            return self.__egg_decls_thunk__()
+        # Catch attribute error, so that it isn't bubbled up as a missing attribute and fallbacks on `__getattr__`
+        # instead raise explicitly
+        except AttributeError as err:
+            msg = "Error resolving declerations"
+            raise RuntimeError(msg) from err
 
 
 @runtime_checkable
@@ -113,6 +121,12 @@ class Declarations:
     _classes: dict[str, ClassDecl] = field(default_factory=dict)
     _rulesets: dict[str, RulesetDecl | CombinedRulesetDecl] = field(default_factory=lambda: {"": RulesetDecl([])})
 
+    @property
+    def default_ruleset(self) -> RulesetDecl:
+        ruleset = self._rulesets[""]
+        assert isinstance(ruleset, RulesetDecl)
+        return ruleset
+
     @classmethod
     def create(cls, *others: DeclerationsLike) -> Declarations:
         others = upcast_declerations(others)
@@ -127,7 +141,7 @@ class Declarations:
 
     def copy(self) -> Declarations:
         new = Declarations()
-        new |= self
+        self.update_other(new)
         return new
 
     def update(self, *others: DeclerationsLike) -> None:
@@ -154,9 +168,13 @@ class Declarations:
         other._functions |= self._functions
         other._classes |= self._classes
         other._constants |= self._constants
+        # Must combine rulesets bc the empty ruleset might be different, bc DefaultRewriteDecl
+        # is added to functions.
+        combined_default_rules: set[RewriteOrRuleDecl] = {*self.default_ruleset.rules, *other.default_ruleset.rules}
         other._rulesets |= self._rulesets
+        other._rulesets[""] = RulesetDecl(list(combined_default_rules))
 
-    def get_callable_decl(self, ref: CallableRef) -> CallableDecl:
+    def get_callable_decl(self, ref: CallableRef) -> CallableDecl:  # noqa: PLR0911
         match ref:
             case FunctionRef(name):
                 return self._functions[name]
@@ -170,8 +188,29 @@ class Declarations:
                 return self._classes[class_name].class_methods[name]
             case PropertyRef(class_name, property_name):
                 return self._classes[class_name].properties[property_name]
+            case InitRef(class_name):
+                init_fn = self._classes[class_name].init
+                assert init_fn
+                return init_fn
         assert_never(ref)
 
+    def set_function_decl(
+        self, ref: FunctionRef | MethodRef | ClassMethodRef | PropertyRef | InitRef, decl: FunctionDecl
+    ) -> None:
+        match ref:
+            case FunctionRef(name):
+                self._functions[name] = decl
+            case MethodRef(class_name, method_name):
+                self._classes[class_name].methods[method_name] = decl
+            case ClassMethodRef(class_name, name):
+                self._classes[class_name].class_methods[name] = decl
+            case PropertyRef(class_name, property_name):
+                self._classes[class_name].properties[property_name] = decl
+            case InitRef(class_name):
+                self._classes[class_name].init = decl
+            case _:
+                assert_never(ref)
+
     def has_method(self, class_name: str, method_name: str) -> bool | None:
         """
         Returns whether the given class has the given method, or None if we cant find the class.
@@ -183,12 +222,20 @@ class Declarations:
     def get_class_decl(self, name: str) -> ClassDecl:
         return self._classes[name]
 
+    def get_paramaterized_class(self, name: str) -> TypeRefWithVars:
+        """
+        Returns a class reference with type parameters, if the class is paramaterized.
+        """
+        type_vars = self._classes[name].type_vars
+        return TypeRefWithVars(name, tuple(map(ClassTypeVarRef, type_vars)))
+
 
 @dataclass
 class ClassDecl:
     egg_name: str | None = None
     type_vars: tuple[str, ...] = ()
     builtin: bool = False
+    init: FunctionDecl | None = None
     class_methods: dict[str, FunctionDecl] = field(default_factory=dict)
     # These have to be seperate from class_methods so that printing them can be done easily
     class_variables: dict[str, ConstantDecl] = field(default_factory=dict)
@@ -293,6 +340,11 @@ class ClassMethodRef:
     method_name: str
 
 
+@dataclass(frozen=True)
+class InitRef:
+    class_name: str
+
+
 @dataclass(frozen=True)
 class ClassVariableRef:
     class_name: str
@@ -305,7 +357,9 @@ class PropertyRef:
     property_name: str
 
 
-CallableRef: TypeAlias = FunctionRef | ConstantRef | MethodRef | ClassMethodRef | ClassVariableRef | PropertyRef
+CallableRef: TypeAlias = (
+    FunctionRef | ConstantRef | MethodRef | ClassMethodRef | InitRef | ClassVariableRef | PropertyRef
+)
 
 
 ##
@@ -378,7 +432,6 @@ class FunctionSignature:
 @dataclass(frozen=True)
 class FunctionDecl:
     signature: FunctionSignature | SpecialFunctions = field(default_factory=FunctionSignature)
-
     # Egg params
     builtin: bool = False
     egg_name: str | None = None
@@ -458,7 +511,7 @@ class CallDecl:
     bound_tp_params: tuple[JustTypeRef, ...] | None = None
 
     def __post_init__(self) -> None:
-        if self.bound_tp_params and not isinstance(self.callable, ClassMethodRef):
+        if self.bound_tp_params and not isinstance(self.callable, ClassMethodRef | InitRef):
             msg = "Cannot bind type parameters to a non-class method callable."
             raise ValueError(msg)
 
@@ -629,7 +682,13 @@ class RuleDecl:
     name: str | None
 
 
-RewriteOrRuleDecl: TypeAlias = RewriteDecl | BiRewriteDecl | RuleDecl
+@dataclass(frozen=True)
+class DefaultRewriteDecl:
+    ref: CallableRef
+    expr: ExprDecl
+
+
+RewriteOrRuleDecl: TypeAlias = RewriteDecl | BiRewriteDecl | RuleDecl | DefaultRewriteDecl
 
 
 @dataclass(frozen=True)

egglog/egraph.py

@@ -352,7 +352,7 @@ class _BaseModule:
         This is the same as defining a nullary function with a high cost.
         # TODO: Rename as declare to match eggglog?
         """
-        return constant(name, tp, egg_name)
+        return constant(name, tp, egg_name=egg_name)
 
     def register(
         self,
@@ -452,6 +452,7 @@ class _ExprMetaclass(type):
         namespace: dict[str, Any],
         egg_sort: str | None = None,
         builtin: bool = False,
+        ruleset: Ruleset | None = None,
     ) -> RuntimeClass | type:
         # If this is the Expr subclass, just return the class
         if not bases:
@@ -467,7 +468,14 @@ class _ExprMetaclass(type):
         # Otherwise, f_locals returns a copy
         # https://peps.python.org/pep-0667/
         decls_thunk = Thunk.fn(
-            _generate_class_decls, namespace, prev_frame, builtin, egg_sort, name, fallback=Declarations
+            _generate_class_decls,
+            namespace,
+            prev_frame,
+            builtin,
+            egg_sort,
+            name,
+            ruleset,
+            fallback=Declarations,
         )
         return RuntimeClass(decls_thunk, TypeRefWithVars(name))
 
@@ -475,8 +483,13 @@ class _ExprMetaclass(type):
         return isinstance(instance, RuntimeExpr)
 
 
-def _generate_class_decls(  # noqa: C901
-    namespace: dict[str, Any], frame: FrameType, builtin: bool, egg_sort: str | None, cls_name: str
+def _generate_class_decls(
+    namespace: dict[str, Any],
+    frame: FrameType,
+    builtin: bool,
+    egg_sort: str | None,
+    cls_name: str,
+    ruleset: Ruleset | None,
 ) -> Declarations:
     """
     Lazy constructor for class declerations to support classes with methods whose types are not yet defined.
@@ -498,9 +511,9 @@ def _generate_class_decls(  # noqa: C901
     for k, v in get_type_hints(_Dummytype, globalns=frame.f_globals, localns=frame.f_locals).items():
         if getattr(v, "__origin__", None) == ClassVar:
             (inner_tp,) = v.__args__
-            type_ref = resolve_type_annotation(decls, inner_tp).to_just()
-            cls_decl.class_variables[k] = ConstantDecl(type_ref)
-
+            type_ref = resolve_type_annotation(decls, inner_tp)
+            cls_decl.class_variables[k] = ConstantDecl(type_ref.to_just())
+            _add_default_rewrite(decls, ClassVariableRef(cls_name, k), type_ref, namespace.pop(k, None), ruleset)
         else:
             msg = f"On class {cls_name}, for attribute '{k}', expected a ClassVar, but got {v}"
             raise NotImplementedError(msg)
@@ -510,13 +523,15 @@ def _generate_class_decls(  # noqa: C901
     ##
 
     # The type ref of self is paramterized by the type vars
-    slf_type_ref = TypeRefWithVars(cls_name, tuple(map(ClassTypeVarRef, type_vars)))
+    TypeRefWithVars(cls_name, tuple(map(ClassTypeVarRef, type_vars)))
 
     # Get all the methods from the class
     filtered_namespace: list[tuple[str, Any]] = [
         (k, v) for k, v in namespace.items() if k not in IGNORED_ATTRIBUTES or isinstance(v, _WrappedMethod)
     ]
 
+    # all methods we should try adding default functions for
+    default_function_refs: dict[ClassMethodRef | MethodRef | PropertyRef, Callable] = {}
     # Then register each of its methods
     for method_name, method in filtered_namespace:
         is_init = method_name == "__init__"
@@ -536,43 +551,26 @@ def _generate_class_decls(  # noqa: C901
             continue
         locals = frame.f_locals
 
-        def create_decl(fn: object, first: Literal["cls"] | TypeRefWithVars) -> FunctionDecl:
-            special_function_name: SpecialFunctions | None = (
-                "fn-partial" if egg_fn == "unstable-fn" else "fn-app" if egg_fn == "unstable-app" else None  # noqa: B023
-            )
-            if special_function_name:
-                return FunctionDecl(
-                    special_function_name,
-                    builtin=True,
-                    egg_name=egg_fn,  # noqa: B023
-                )
-
-            return _fn_decl(
-                decls,
-                egg_fn,  # noqa: B023
-                fn,
-                locals,  # noqa: B023
-                default,  # noqa: B023
-                cost,  # noqa: B023
-                merge,  # noqa: B023
-                on_merge,  # noqa: B023
-                mutates,  # noqa: B023
-                builtin,
-                first,
-                is_init,  # noqa: B023
-                unextractable,  # noqa: B023
-            )
-
+        ref: ClassMethodRef | MethodRef | PropertyRef | InitRef
         match fn:
             case classmethod():
-                cls_decl.class_methods[method_name] = create_decl(fn.__func__, "cls")
+                ref = ClassMethodRef(cls_name, method_name)
+                fn = fn.__func__
             case property():
-                cls_decl.properties[method_name] = create_decl(fn.fget, slf_type_ref)
+                ref = PropertyRef(cls_name, method_name)
+                fn = fn.fget
             case _:
-                if is_init:
-                    cls_decl.class_methods[method_name] = create_decl(fn, slf_type_ref)
-                else:
-                    cls_decl.methods[method_name] = create_decl(fn, slf_type_ref)
+                ref = InitRef(cls_name) if is_init else MethodRef(cls_name, method_name)
+
+        _fn_decl(decls, egg_fn, ref, fn, locals, default, cost, merge, on_merge, mutates, builtin, unextractable)
+
+        if not builtin and not isinstance(ref, InitRef) and not mutates:
+            default_function_refs[ref] = fn
+
+    # Add all rewrite methods at the end so that all methods are registered first and can be accessed
+    # in the bodies
+    for ref, fn in default_function_refs.items():
+        _add_default_rewrite_function(decls, ref, fn, ruleset)
 
     return decls
 
@@ -591,6 +589,7 @@ def function(
     mutates_first_arg: bool = False,
     unextractable: bool = False,
     builtin: bool = False,
+    ruleset: Ruleset | None = None,
 ) -> Callable[[CALLABLE], CALLABLE]: ...
 
 
@@ -604,6 +603,7 @@ def function(
     on_merge: Callable[[EXPR, EXPR], Iterable[ActionLike]] | None = None,
     mutates_first_arg: bool = False,
     unextractable: bool = False,
+    ruleset: Ruleset | None = None,
 ) -> Callable[[Callable[P, EXPR]], Callable[P, EXPR]]: ...
 
 
@@ -634,15 +634,17 @@ class _FunctionConstructor:
     merge: Callable[[RuntimeExpr, RuntimeExpr], RuntimeExpr] | None = None
     on_merge: Callable[[RuntimeExpr, RuntimeExpr], Iterable[ActionLike]] | None = None
     unextractable: bool = False
+    ruleset: Ruleset | None = None
 
     def __call__(self, fn: Callable[..., RuntimeExpr]) -> RuntimeFunction:
         return RuntimeFunction(Thunk.fn(self.create_decls, fn), FunctionRef(fn.__name__))
 
     def create_decls(self, fn: Callable[..., RuntimeExpr]) -> Declarations:
         decls = Declarations()
-        decls._functions[fn.__name__] = _fn_decl(
+        _fn_decl(
             decls,
             self.egg_fn,
+            (ref := FunctionRef(fn.__name__)),
             fn,
             self.hint_locals,
             self.default,
@@ -653,12 +655,14 @@ class _FunctionConstructor:
             self.builtin,
             unextractable=self.unextractable,
         )
+        _add_default_rewrite_function(decls, ref, fn, self.ruleset)
         return decls
 
 
 def _fn_decl(
     decls: Declarations,
     egg_name: str | 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)
@@ -669,14 +673,23 @@ def _fn_decl(
     on_merge: Callable[[RuntimeExpr, RuntimeExpr], Iterable[ActionLike]] | None,
     mutates_first_arg: bool,
     is_builtin: bool,
-    # The first arg is either cls, for a classmethod, a self type, or none for a function
-    first_arg: Literal["cls"] | TypeOrVarRef | None = None,
-    is_init: bool = False,
     unextractable: bool = False,
-) -> FunctionDecl:
+) -> None:
+    """
+    Sets the function decl for the function object.
+    """
     if not isinstance(fn, FunctionType):
         raise NotImplementedError(f"Can only generate function decls for functions not {fn}  {type(fn)}")
 
+    # partial function creation and calling are handled with a special case in the type checker, so don't
+    # use the normal logic
+    special_function_name: SpecialFunctions | None = (
+        "fn-partial" if egg_name == "unstable-fn" else "fn-app" if egg_name == "unstable-app" else None
+    )
+    if special_function_name:
+        decls.set_function_decl(ref, FunctionDecl(special_function_name, builtin=True, egg_name=egg_name))
+        return
+
     hint_globals = fn.__globals__.copy()
     # Copy Callable into global if not present bc sometimes it gets automatically removed by ruff to type only block
     # https://docs.astral.sh/ruff/rules/typing-only-standard-library-import/
@@ -687,8 +700,8 @@ def _fn_decl(
 
     params = list(signature(fn).parameters.values())
 
-    # If this is an init function, or a classmethod, remove the first arg name
-    if is_init or first_arg == "cls":
+    # If this is an init function, or a classmethod, the first arg is not used
+    if isinstance(ref, ClassMethodRef | InitRef):
         params = params[1:]
 
     if _last_param_variable(params):
@@ -698,9 +711,8 @@ def _fn_decl(
     else:
         var_arg_type = None
     arg_types = tuple(
-        first_arg
-        # If the first arg is a self, and this not an __init__ fn, add this as a typeref
-        if i == 0 and isinstance(first_arg, ClassTypeVarRef | TypeRefWithVars) and not is_init
+        decls.get_paramaterized_class(ref.class_name)
+        if i == 0 and isinstance(ref, MethodRef | PropertyRef)
         else resolve_type_annotation(decls, hints[t.name])
         for i, t in enumerate(params)
     )
@@ -713,17 +725,15 @@ def _fn_decl(
 
     decls.update(*arg_defaults)
 
-    # If this is an init fn use the first arg as the return type
-    if is_init:
-        assert not mutates_first_arg
-        if not isinstance(first_arg, ClassTypeVarRef | TypeRefWithVars):
-            msg = "Init function must have a self type"
-            raise ValueError(msg)
-        return_type = first_arg
-    elif mutates_first_arg:
-        return_type = arg_types[0]
-    else:
-        return_type = resolve_type_annotation(decls, hints["return"])
+    return_type = (
+        decls.get_paramaterized_class(ref.class_name)
+        if isinstance(ref, InitRef)
+        else arg_types[0]
+        if mutates_first_arg
+        else resolve_type_annotation(decls, hints["return"])
+    )
+
+    arg_names = tuple(t.name for t in params)
 
     decls |= default
     merged = (
@@ -747,12 +757,12 @@ def _fn_decl(
         )
     )
     decls.update(*merge_action)
-    return FunctionDecl(
+    decl = FunctionDecl(
         FunctionSignature(
             return_type=None if mutates_first_arg else return_type,
             var_arg_type=var_arg_type,
             arg_types=arg_types,
-            arg_names=tuple(t.name for t in params),
+            arg_names=arg_names,
             arg_defaults=tuple(a.__egg_typed_expr__.expr if a is not None else None for a in arg_defaults),
         ),
         cost=cost,
@@ -763,6 +773,7 @@ def _fn_decl(
         default=None if default is None else default.__egg_typed_expr__.expr,
         on_merge=tuple(a.action for a in merge_action),
     )
+    decls.set_function_decl(ref, decl)
 
 
 # Overload to support aritys 0-4 until variadic generic support map, so we can map from type to value
@@ -804,19 +815,73 @@ def _relation_decls(name: str, tps: tuple[type, ...], egg_fn: str | None) -> Dec
     return decls
 
 
-def constant(name: str, tp: type[EXPR], egg_name: str | None = None) -> EXPR:
+def constant(
+    name: str,
+    tp: type[EXPR],
+    default_replacement: EXPR | None = None,
+    /,
+    *,
+    egg_name: str | None = None,
+    ruleset: Ruleset | None = None,
+) -> EXPR:
     """
     A "constant" is implemented as the instantiation of a value that takes no args.
     This creates a function with `name` and return type `tp` and returns a value of it being called.
     """
-    return cast(EXPR, RuntimeExpr(Thunk.fn(_constant_thunk, name, tp, egg_name)))
+    return cast(EXPR, RuntimeExpr(Thunk.fn(_constant_thunk, name, tp, egg_name, default_replacement, ruleset)))
 
 
-def _constant_thunk(name: str, tp: type, egg_name: str | None) -> tuple[Declarations, TypedExprDecl]:
+def _constant_thunk(
+    name: str, tp: type, egg_name: str | None, default_replacement: object, ruleset: Ruleset | None
+) -> tuple[Declarations, TypedExprDecl]:
     decls = Declarations()
-    type_ref = resolve_type_annotation(decls, tp).to_just()
-    decls._constants[name] = ConstantDecl(type_ref, egg_name)
-    return decls, TypedExprDecl(type_ref, CallDecl(ConstantRef(name)))
+    type_ref = resolve_type_annotation(decls, tp)
+    callable_ref = ConstantRef(name)
+    decls._constants[name] = ConstantDecl(type_ref.to_just(), egg_name)
+    _add_default_rewrite(decls, callable_ref, type_ref, default_replacement, ruleset)
+    return decls, TypedExprDecl(type_ref.to_just(), CallDecl(callable_ref))
+
+
+def _add_default_rewrite_function(decls: Declarations, ref: CallableRef, fn: Callable, ruleset: Ruleset | None) -> None:
+    """
+    Adds a default rewrite for a function, by calling the functions with vars and adding it if it is not None.
+    """
+    callable_decl = decls.get_callable_decl(ref)
+    assert isinstance(callable_decl, FunctionDecl)
+    signature = callable_decl.signature
+    assert isinstance(signature, FunctionSignature)
+
+    var_args: list[object] = [
+        RuntimeExpr.__from_value__(decls, TypedExprDecl(tp.to_just(), VarDecl(_rule_var_name(name))))
+        for name, tp in zip(signature.arg_names, signature.arg_types, strict=False)
+    ]
+    # If this is a classmethod, add the class as the first arg
+    if isinstance(ref, ClassMethodRef):
+        tp = decls.get_paramaterized_class(ref.class_name)
+        var_args.insert(0, RuntimeClass(Thunk.value(decls), tp))
+    _add_default_rewrite(decls, ref, signature.semantic_return_type, fn(*var_args), ruleset)
+
+
+def _add_default_rewrite(
+    decls: Declarations, ref: CallableRef, type_ref: TypeOrVarRef, default_rewrite: object, ruleset: Ruleset | None
+) -> None:
+    """
+    Adds a default rewrite for the callable, if the default rewrite is not None
+
+    Will add it to the ruleset if it is passed in, or add it to the default ruleset on the passed in decls if not.
+    """
+    if default_rewrite is None:
+        return
+    resolved_value = resolve_literal(type_ref, default_rewrite)
+    rewrite_decl = DefaultRewriteDecl(ref, resolved_value.__egg_typed_expr__.expr)
+    if ruleset:
+        ruleset_decls = ruleset._current_egg_decls
+        ruleset_decl = ruleset.__egg_ruleset__
+    else:
+        ruleset_decls = decls
+        ruleset_decl = decls.default_ruleset
+    ruleset_decl.rules.append(rewrite_decl)
+    ruleset_decls |= resolved_value
 
 
 def _last_param_variable(params: list[Parameter]) -> bool:
@@ -991,6 +1056,7 @@ class EGraph(_BaseModule):
         action = let(name, expr)
         self.register(action)
         runtime_expr = to_runtime_expr(expr)
+        self._add_decls(runtime_expr)
         return cast(
             EXPR,
             RuntimeExpr.__from_value__(
@@ -1016,7 +1082,8 @@ class EGraph(_BaseModule):
         self._add_decls(runtime_expr, schedule)
         egg_schedule = self._state.schedule_to_egg(schedule.schedule)
         typed_expr = runtime_expr.__egg_typed_expr__
-        egg_expr = self._state.expr_to_egg(typed_expr.expr)
+        # Must also register type
+        egg_expr = self._state.typed_expr_to_egg(typed_expr)
         self._egraph.run_program(bindings.Simplify(egg_expr, egg_schedule))
         extract_report = self._egraph.extract_report()
         if not isinstance(extract_report, bindings.Best):
@@ -1121,8 +1188,7 @@ class EGraph(_BaseModule):
         return [cast(EXPR, RuntimeExpr.__from_value__(self.__egg_decls__, expr)) for expr in new_exprs]
 
     def _run_extract(self, typed_expr: TypedExprDecl, n: int) -> bindings._ExtractReport:
-        self._state.type_ref_to_egg(typed_expr.tp)
-        expr = self._state.expr_to_egg(typed_expr.expr)
+        expr = self._state.typed_expr_to_egg(typed_expr)
         self._egraph.run_program(bindings.ActionCommand(bindings.Extract(expr, bindings.Lit(bindings.Int(n)))))
         extract_report = self._egraph.extract_report()
         if not extract_report:
@@ -1181,7 +1247,7 @@ class EGraph(_BaseModule):
         runtime_expr = to_runtime_expr(expr)
         self._add_decls(runtime_expr)
         typed_expr = runtime_expr.__egg_typed_expr__
-        egg_expr = self._state.expr_to_egg(typed_expr.expr)
+        egg_expr = self._state.typed_expr_to_egg(typed_expr)
         match typed_expr.tp:
             case JustTypeRef("i64"):
                 return self._egraph.eval_i64(egg_expr)
@@ -1878,7 +1944,7 @@ def _rewrite_or_rule_generator(gen: RewriteOrRuleGenerator, frame: FrameType) ->
     if "Callable" not in globals:
         globals["Callable"] = Callable
     hints = get_type_hints(gen, globals, frame.f_locals)
-    args = [_var(p.name, hints[p.name]) for p in signature(gen).parameters.values()]
+    args = [_var(_rule_var_name(p.name), hints[p.name]) for p in signature(gen).parameters.values()]
     return list(gen(*args))  # type: ignore[misc]
 
 

egglog/egraph_state.py

@@ -19,7 +19,7 @@ from .type_constraint_solver import TypeConstraintError, TypeConstraintSolver
 if TYPE_CHECKING:
     from collections.abc import Iterable
 
-__all__ = ["EGraphState", "GLOBAL_PY_OBJECT_SORT"]
+__all__ = ["EGraphState", "GLOBAL_PY_OBJECT_SORT", "_rule_var_name"]
 
 # Create a global sort for python objects, so we can store them without an e-graph instance
 # Needed when serializing commands to egg commands when creating modules
@@ -98,7 +98,8 @@ class EGraphState:
                 for rule in rules:
                     if rule in added_rules:
                         continue
-                    self.egraph.run_program(self.command_to_egg(rule, name))
+                    cmd = self.command_to_egg(rule, name)
+                    self.egraph.run_program(cmd)
                     added_rules.add(rule)
             case CombinedRulesetDecl(rulesets):
                 if name in self.rulesets:
@@ -115,8 +116,8 @@ class EGraphState:
             case RewriteDecl(tp, lhs, rhs, conditions) | BiRewriteDecl(tp, lhs, rhs, conditions):
                 self.type_ref_to_egg(tp)
                 rewrite = bindings.Rewrite(
-                    self.expr_to_egg(lhs),
-                    self.expr_to_egg(rhs),
+                    self._expr_to_egg(lhs),
+                    self._expr_to_egg(rhs),
                     [self.fact_to_egg(c) for c in conditions],
                 )
                 return (
@@ -130,6 +131,16 @@ class EGraphState:
                     [self.fact_to_egg(f) for f in body],
                 )
                 return bindings.RuleCommand(name or "", ruleset, rule)
+            case DefaultRewriteDecl(ref, expr):
+                decl = self.__egg_decls__.get_callable_decl(ref).to_function_decl()
+                sig = decl.signature
+                assert isinstance(sig, FunctionSignature)
+                args = tuple(
+                    TypedExprDecl(tp.to_just(), VarDecl(_rule_var_name(name)))
+                    for name, tp in zip(sig.arg_names, sig.arg_types, strict=False)
+                )
+                rewrite_decl = RewriteDecl(sig.semantic_return_type.to_just(), CallDecl(ref, args), expr, (), False)
+                return self.command_to_egg(rewrite_decl, ruleset)
             case _:
                 assert_never(cmd)
 
@@ -139,13 +150,13 @@ class EGraphState:
                 return bindings.Let(name, self.typed_expr_to_egg(typed_expr))
             case SetDecl(tp, call, rhs):
                 self.type_ref_to_egg(tp)
-                call_ = self.expr_to_egg(call)
-                return bindings.Set(call_.name, call_.args, self.expr_to_egg(rhs))
+                call_ = self._expr_to_egg(call)
+                return bindings.Set(call_.name, call_.args, self._expr_to_egg(rhs))
             case ExprActionDecl(typed_expr):
                 return bindings.Expr_(self.typed_expr_to_egg(typed_expr))
             case ChangeDecl(tp, call, change):
                 self.type_ref_to_egg(tp)
-                call_ = self.expr_to_egg(call)
+                call_ = self._expr_to_egg(call)
                 egg_change: bindings._Change
                 match change:
                     case "delete":
@@ -157,7 +168,7 @@ class EGraphState:
                 return bindings.Change(egg_change, call_.name, call_.args)
             case UnionDecl(tp, lhs, rhs):
                 self.type_ref_to_egg(tp)
-                return bindings.Union(self.expr_to_egg(lhs), self.expr_to_egg(rhs))
+                return bindings.Union(self._expr_to_egg(lhs), self._expr_to_egg(rhs))
             case PanicDecl(name):
                 return bindings.Panic(name)
             case _:
@@ -167,7 +178,7 @@ class EGraphState:
         match fact:
             case EqDecl(tp, exprs):
                 self.type_ref_to_egg(tp)
-                return bindings.Eq([self.expr_to_egg(e) for e in exprs])
+                return bindings.Eq([self._expr_to_egg(e) for e in exprs])
             case ExprFactDecl(typed_expr):
                 return bindings.Fact(self.typed_expr_to_egg(typed_expr))
             case _:
@@ -201,8 +212,8 @@ class EGraphState:
                             [self.type_ref_to_egg(a.to_just()) for a in signature.arg_types],
                             self.type_ref_to_egg(signature.semantic_return_type.to_just()),
                         ),
-                        self.expr_to_egg(decl.default) if decl.default else None,
-                        self.expr_to_egg(decl.merge) if decl.merge else None,
+                        self._expr_to_egg(decl.default) if decl.default else None,
+                        self._expr_to_egg(decl.merge) if decl.merge else None,
                         [self.action_to_egg(a) for a in decl.on_merge],
                         decl.cost,
                         decl.unextractable,
@@ -245,6 +256,8 @@ class EGraphState:
         if decl.builtin:
             for method in decl.class_methods:
                 self.callable_ref_to_egg(ClassMethodRef(ref.name, method))
+            if decl.init:
+                self.callable_ref_to_egg(InitRef(ref.name))
 
         return egg_name
 
@@ -261,15 +274,15 @@ class EGraphState:
 
     def typed_expr_to_egg(self, typed_expr_decl: TypedExprDecl) -> bindings._Expr:
         self.type_ref_to_egg(typed_expr_decl.tp)
-        return self.expr_to_egg(typed_expr_decl.expr)
+        return self._expr_to_egg(typed_expr_decl.expr)
 
     @overload
-    def expr_to_egg(self, expr_decl: CallDecl) -> bindings.Call: ...
+    def _expr_to_egg(self, expr_decl: CallDecl) -> bindings.Call: ...
 
     @overload
-    def expr_to_egg(self, expr_decl: ExprDecl) -> bindings._Expr: ...
+    def _expr_to_egg(self, expr_decl: ExprDecl) -> bindings._Expr: ...
 
-    def expr_to_egg(self, expr_decl: ExprDecl) -> bindings._Expr:
+    def _expr_to_egg(self, expr_decl: ExprDecl) -> bindings._Expr:
         """
         Convert an ExprDecl to an egg expression.
 
@@ -307,7 +320,7 @@ class EGraphState:
             case PyObjectDecl(value):
                 res = GLOBAL_PY_OBJECT_SORT.store(value)
             case PartialCallDecl(call_decl):
-                egg_fn_call = self.expr_to_egg(call_decl)
+                egg_fn_call = self._expr_to_egg(call_decl)
                 res = bindings.Call("unstable-fn", [bindings.Lit(bindings.String(egg_fn_call.name)), *egg_fn_call.args])
             case _:
                 assert_never(expr_decl.expr)
@@ -355,6 +368,8 @@ def _generate_callable_egg_name(ref: CallableRef) -> str:
             | PropertyRef(cls_name, name)
         ):
             return f"{cls_name}_{name}"
+        case InitRef(cls_name):
+            return f"{cls_name}___init__"
         case _:
             assert_never(ref)
 
@@ -427,8 +442,11 @@ class FromEggState:
             possible_types: Iterable[JustTypeRef | None]
             signature = self.decls.get_callable_decl(callable_ref).to_function_decl().signature
             assert isinstance(signature, FunctionSignature)
-            if isinstance(callable_ref, ClassMethodRef):
-                possible_types = self.state._get_possible_types(callable_ref.class_name)
+            if isinstance(callable_ref, ClassMethodRef | InitRef | MethodRef):
+                # Need OR in case we have class method whose class whas never added as a sort, which would happen
+                # if the class method didn't return that type and no other function did. In this case, we don't need
+                # to care about the type vars and we we don't need to bind any possible type.
+                possible_types = self.state._get_possible_types(callable_ref.class_name) or [None]
                 cls_name = callable_ref.class_name
             else:
                 possible_types = [None]
@@ -437,7 +455,6 @@ class FromEggState:
                 tcs = TypeConstraintSolver(self.decls)
                 if possible_type and possible_type.args:
                     tcs.bind_class(possible_type)
-
                 try:
                     arg_types, bound_tp_params = tcs.infer_arg_types(
                         signature.arg_types, signature.semantic_return_type, signature.var_arg_type, tp, cls_name
@@ -445,7 +462,14 @@ class FromEggState:
                 except TypeConstraintError:
                     continue
                 args = tuple(self.resolve_term(a, tp) for a, tp in zip(term.args, arg_types, strict=False))
-                return CallDecl(callable_ref, args, bound_tp_params)
+
+                return CallDecl(
+                    callable_ref,
+                    args,
+                    # Don't include bound type params if this is just a method, we only needed them for type resolution
+                    # 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}")
 
     def resolve_term(self, term_id: int, tp: JustTypeRef) -> TypedExprDecl:
@@ -454,3 +478,10 @@ class FromEggState:
         except KeyError:
             res = self.cache[term_id] = self.from_expr(tp, self.termdag.nodes[term_id])
             return res
+
+
+def _rule_var_name(s: str) -> str:
+    """
+    Create a hidden variable name, for rewrites, so that let bindings or function won't conflict with it
+    """
+    return f"__var__{s}"

egglog/examples/higher_order_functions.py

@@ -0,0 +1,50 @@
+# mypy: disable-error-code="empty-body"
+"""
+Higher Order Functions
+======================
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from egglog import *
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+class Math(Expr):
+    def __init__(self, i: i64Like) -> None: ...
+
+    def __add__(self, other: Math) -> Math: ...
+
+
+class MathList(Expr):
+    def __init__(self) -> None: ...
+
+    def append(self, i: Math) -> MathList: ...
+
+    def map(self, f: Callable[[Math], Math]) -> MathList: ...
+
+
+@ruleset
+def math_ruleset(i: i64, j: i64, xs: MathList, x: Math, f: Callable[[Math], Math]):  # noqa: ANN201
+    yield rewrite(Math(i) + Math(j)).to(Math(i + j))
+    yield rewrite(xs.append(x).map(f)).to(xs.map(f).append(f(x)))
+    yield rewrite(MathList().map(f)).to(MathList())
+
+
+@function(ruleset=math_ruleset)
+def increment_by_one(x: Math) -> Math:
+    return x + Math(1)
+
+
+egraph = EGraph()
+x = egraph.let("x", MathList().append(Math(1)).append(Math(2)))
+y = egraph.let("y", x.map(increment_by_one))
+egraph.run(math_ruleset.saturate())
+
+egraph.check(eq(y).to(MathList().append(Math(2)).append(Math(3))))
+
+egraph

egglog/exp/array_api.py

@@ -882,7 +882,10 @@ converter(IntOrTuple, OptionalIntOrTuple, OptionalIntOrTuple.some)
 
 @function
 def asarray(
-    a: NDArray, dtype: OptionalDType = OptionalDType.none, copy: OptionalBool = OptionalBool.none
+    a: NDArray,
+    dtype: OptionalDType = OptionalDType.none,
+    copy: OptionalBool = OptionalBool.none,
+    device: OptionalDevice = OptionalDevice.none,
 ) -> NDArray: ...
 
 

egglog/pretty.py

@@ -166,6 +166,8 @@ class TraverseContext:
                 pass
             case EqDecl(_, decls) | SequenceDecl(decls) | RulesetDecl(decls):
                 for de in decls:
+                    if isinstance(de, DefaultRewriteDecl):
+                        continue
                     self(de)
             case CallDecl(_, exprs, _):
                 for e in exprs:
@@ -178,6 +180,8 @@ class TraverseContext:
                 self(c)
             case CombinedRulesetDecl(_):
                 pass
+            case DefaultRewriteDecl():
+                pass
             case _:
                 assert_never(decl)
 
@@ -276,7 +280,7 @@ class PrettyContext:
             case RulesetDecl(rules):
                 if ruleset_name:
                     return f"ruleset(name={ruleset_name!r})", f"ruleset_{ruleset_name}"
-                args = ", ".join(map(self, rules))
+                args = ", ".join(self(r) for r in rules if not isinstance(r, DefaultRewriteDecl))
                 return f"ruleset({args})", "ruleset"
             case CombinedRulesetDecl(rulesets):
                 if ruleset_name:
@@ -298,6 +302,9 @@ class PrettyContext:
                     return ruleset_str, "schedule"
                 args = ", ".join(map(self, until))
                 return f"run({ruleset_str}, {args})", "schedule"
+            case DefaultRewriteDecl():
+                msg = "default rewrites should not be pretty printed"
+                raise TypeError(msg)
         assert_never(decl)
 
     def _call(
@@ -370,10 +377,8 @@ class PrettyContext:
             case FunctionRef(name):
                 return name, args
             case ClassMethodRef(class_name, method_name):
-                fn_str = str(JustTypeRef(class_name, bound_tp_params or ()))
-                if method_name != "__init__":
-                    fn_str += f".{method_name}"
-                return fn_str, args
+                tp_ref = JustTypeRef(class_name, bound_tp_params or ())
+                return f"{tp_ref}.{method_name}", args
             case MethodRef(_class_name, method_name):
                 slf, *args = args
                 slf = self(slf, parens=True)
@@ -400,6 +405,9 @@ class PrettyContext:
                 return f"{class_name}.{variable_name}"
             case PropertyRef(_class_name, property_name):
                 return f"{self(args[0], parens=True)}.{property_name}"
+            case InitRef(class_name):
+                tp_ref = JustTypeRef(class_name, bound_tp_params or ())
+                return str(tp_ref), args
         assert_never(ref)
 
     def _generate_name(self, typ: str) -> str:
@@ -434,6 +442,8 @@ def _pretty_callable(ref: CallableRef) -> str:
             | PropertyRef(class_name, method_name)
         ):
             return f"{class_name}.{method_name}"
+        case InitRef(class_name):
+            return class_name
         case ConstantRef(_):
             msg = "Constants should not be callable"
             raise NotImplementedError(msg)

egglog/runtime.py

@@ -163,10 +163,8 @@ class RuntimeClass(DelayedDeclerations):
             return RuntimeExpr.__from_value__(
                 self.__egg_decls__, TypedExprDecl(self.__egg_tp__.to_just(), LitDecl(None))
             )
-
-        return RuntimeFunction(
-            Thunk.value(self.__egg_decls__), ClassMethodRef(name, "__init__"), self.__egg_tp__.to_just()
-        )(*args, **kwargs)  # type: ignore[arg-type]
+        fn = RuntimeFunction(Thunk.value(self.__egg_decls__), InitRef(name), self.__egg_tp__.to_just())
+        return fn(*args, **kwargs)  # type: ignore[arg-type]
 
     def __dir__(self) -> list[str]:
         cls_decl = self.__egg_decls__.get_class_decl(self.__egg_tp__.name)
@@ -277,7 +275,10 @@ class RuntimeFunction(DelayedDeclerations):
 
         # Turn all keyword args into positional args
         py_signature = to_py_signature(signature, self.__egg_decls__, _egg_partial_function)
-        bound = py_signature.bind(*args, **kwargs)
+        try:
+            bound = py_signature.bind(*args, **kwargs)
+        except TypeError as err:
+            raise TypeError(f"Failed to call {self} with args {args} and kwargs {kwargs}") from err
         del kwargs
         bound.apply_defaults()
         assert not bound.kwargs
@@ -310,7 +311,9 @@ class RuntimeFunction(DelayedDeclerations):
         return_tp = tcs.infer_return_type(
             signature.arg_types, signature.semantic_return_type, signature.var_arg_type, arg_types, cls_name
         )
-        bound_params = cast(JustTypeRef, bound_tp).args if isinstance(self.__egg_ref__, ClassMethodRef) else None
+        bound_params = (
+            cast(JustTypeRef, bound_tp).args if isinstance(self.__egg_ref__, ClassMethodRef | InitRef) else None
+        )
         # If we were using unstable-app to call a funciton, add that function back as the first arg.
         if function_value:
             arg_exprs = (function_value, *arg_exprs)

egglog/thunk.py

@@ -21,7 +21,7 @@ class Thunk(Generic[T, Unpack[TS]]):
     Cached delayed function call.
     """
 
-    state: Resolved[T] | Unresolved[T, Unpack[TS]] | Resolving[T]
+    state: Resolved[T] | Unresolved[T, Unpack[TS]] | Resolving[T] | Error
 
     @classmethod
     def fn(
@@ -44,14 +44,21 @@ class Thunk(Generic[T, Unpack[TS]]):
                 return value
             case Unresolved(fn, args, fallback):
                 self.state = Resolving(fallback)
-                res = fn(*args)
-                self.state = Resolved(res)
-                return res
+                try:
+                    res = fn(*args)
+                except Exception as e:
+                    self.state = Error(e)
+                    raise
+                else:
+                    self.state = Resolved(res)
+                    return res
             case Resolving(fallback):
                 if fallback is None:
                     msg = "Recursively resolving thunk without fallback"
                     raise ValueError(msg)
                 return fallback()
+            case Error(e):
+                raise e
 
 
 @dataclass
@@ -69,3 +76,8 @@ class Unresolved(Generic[T, Unpack[TS]]):
 @dataclass
 class Resolving(Generic[T]):
     fallback: Callable[[], T] | None
+
+
+@dataclass
+class Error:
+    e: Exception

egglog/type_constraint_solver.py

@@ -79,9 +79,10 @@ class TypeConstraintSolver:
         Also returns the bound type params if the class name is passed in.
         """
         self._infer_typevars(fn_return, return_, cls_name)
-        arg_types = (
-            self._subtitute_typevars(a, cls_name) for a in chain(fn_args, repeat(fn_var_args) if fn_var_args else [])
-        )
+        arg_types: Iterable[JustTypeRef] = [self._subtitute_typevars(a, cls_name) for a in fn_args]
+        if fn_var_args:
+            # Need to be generator so it can be infinite for variable args
+            arg_types = chain(arg_types, repeat(self._subtitute_typevars(fn_var_args, cls_name)))
         bound_typevars = (
             tuple(
                 v
@@ -132,8 +133,8 @@ class TypeConstraintSolver:
     def _subtitute_typevars(self, tp: TypeOrVarRef, cls_name: str | None) -> JustTypeRef:
         match tp:
             case ClassTypeVarRef(name):
+                assert cls_name is not None
                 try:
-                    assert cls_name is not None
                     return self._cls_typevar_index_to_type[cls_name][name]
                 except KeyError as e:
                     raise TypeConstraintError(f"Not enough bound typevars for {tp}") from e

{egglog-7.1.0.dist-info → egglog-7.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: egglog
-Version: 7.1.0
+Version: 7.2.0
 Classifier: Environment :: MacOS X
 Classifier: Environment :: Win32 (MS Windows)
 Classifier: Intended Audience :: Developers
@@ -19,22 +19,6 @@ Classifier: Typing :: Typed
 Requires-Dist: typing-extensions
 Requires-Dist: black
 Requires-Dist: graphviz
-Requires-Dist: scikit-learn; extra == 'array'
-Requires-Dist: array_api_compat; extra == 'array'
-Requires-Dist: numba==0.59.1; extra == 'array'
-Requires-Dist: llvmlite==0.42.0; extra == 'array'
-Requires-Dist: pre-commit; extra == 'dev'
-Requires-Dist: ruff; extra == 'dev'
-Requires-Dist: mypy; extra == 'dev'
-Requires-Dist: anywidget[dev]; extra == 'dev'
-Requires-Dist: egglog[docs,test]; extra == 'dev'
-Requires-Dist: pytest; extra == 'test'
-Requires-Dist: mypy; extra == 'test'
-Requires-Dist: syrupy; extra == 'test'
-Requires-Dist: egglog[array]; extra == 'test'
-Requires-Dist: pytest-codspeed; extra == 'test'
-Requires-Dist: pytest-benchmark; extra == 'test'
-Requires-Dist: pytest-xdist; extra == 'test'
 Requires-Dist: pydata-sphinx-theme; extra == 'docs'
 Requires-Dist: myst-nb; extra == 'docs'
 Requires-Dist: sphinx-autodoc-typehints; extra == 'docs'
@@ -47,10 +31,26 @@ Requires-Dist: egglog[array]; extra == 'docs'
 Requires-Dist: line-profiler; extra == 'docs'
 Requires-Dist: sphinxcontrib-mermaid; extra == 'docs'
 Requires-Dist: ablog; extra == 'docs'
+Requires-Dist: pytest; extra == 'test'
+Requires-Dist: mypy; extra == 'test'
+Requires-Dist: syrupy; extra == 'test'
+Requires-Dist: egglog[array]; extra == 'test'
+Requires-Dist: pytest-codspeed; extra == 'test'
+Requires-Dist: pytest-benchmark; extra == 'test'
+Requires-Dist: pytest-xdist; extra == 'test'
+Requires-Dist: scikit-learn; extra == 'array'
+Requires-Dist: array_api_compat; extra == 'array'
+Requires-Dist: numba==0.59.1; extra == 'array'
+Requires-Dist: llvmlite==0.42.0; extra == 'array'
+Requires-Dist: pre-commit; extra == 'dev'
+Requires-Dist: ruff; extra == 'dev'
+Requires-Dist: mypy; extra == 'dev'
+Requires-Dist: anywidget[dev]; extra == 'dev'
+Requires-Dist: egglog[docs,test]; extra == 'dev'
+Provides-Extra: docs
+Provides-Extra: test
 Provides-Extra: array
 Provides-Extra: dev
-Provides-Extra: test
-Provides-Extra: docs
 License-File: LICENSE
 Summary: e-graphs in Python built around the the egglog rust library
 License: MIT

{egglog-7.1.0.dist-info → egglog-7.2.0.dist-info}/RECORD

@@ -1,16 +1,17 @@
-egglog-7.1.0.dist-info/METADATA,sha256=t8YGDTCeCIpcJMDrAiicXrJAARElS-vM1R6-GY4wNzE,3829
-egglog-7.1.0.dist-info/WHEEL,sha256=keLBtIOE7ZgLyd8Cijw37iqv72a2jhfxMdDu8Unr7zo,96
-egglog-7.1.0.dist-info/license_files/LICENSE,sha256=TfaboMVZ81Q6OUaKjU7z6uVjSlcGKclLYcOpgDbm9_s,1091
+egglog-7.2.0.dist-info/METADATA,sha256=O3Uq5oIIHWZ-bG2XVILykl-0dYyuNuoR9Dxj38v6bXs,3829
+egglog-7.2.0.dist-info/WHEEL,sha256=keLBtIOE7ZgLyd8Cijw37iqv72a2jhfxMdDu8Unr7zo,96
+egglog-7.2.0.dist-info/license_files/LICENSE,sha256=TfaboMVZ81Q6OUaKjU7z6uVjSlcGKclLYcOpgDbm9_s,1091
 egglog/bindings.pyi,sha256=iFdtYHqPjuVt46qh1IE81cOI-lbGgPISKQB-3ERNzhI,11770
 egglog/builtins.py,sha256=p5oZLDleCgQC8G2Zp2EvmqsfVnpgmARqUHqosKSSKnQ,13120
 egglog/config.py,sha256=mALVaxh7zmGrbuyzaVKVmYKcu1lF703QsKJw8AF7gSM,176
 egglog/conversion.py,sha256=4JhocGd1_nwmFMVNCzDDwj6aWfhBFTX2hm_7Xc_DiUM,6328
-egglog/declarations.py,sha256=LFClid1ojrQFF7ezJSanwZA89j01fjJ-NkvH0VUH0Wk,16199
-egglog/egraph.py,sha256=WRu_omk-5i53wokpkSTw9wA2EqW_WI2-ZX8340_w3q8,66198
-egglog/egraph_state.py,sha256=ic6MEL41-DjozBCe-GSKqHM2vkabuC865odeg_zT964,20074
+egglog/declarations.py,sha256=WpZ1yyw4vlWm3Xi_otPP2IU51IqNcDmJA5y20RcBFBQ,18583
+egglog/egraph.py,sha256=Q3ZJhtzzWcsukuarT0vT48-DkfHe17adLLUHn-pILPE,68657
+egglog/egraph_state.py,sha256=kXCNt4zenG1meRv5T6B8xI5oH_NreA06JMhmqv_F-og,21802
 egglog/examples/bool.py,sha256=pWZTjfXR1cFy3KcihLBU5AF5rn83ImORlhUUJ1YiAXc,733
 egglog/examples/eqsat_basic.py,sha256=ORXFYYEDsEZK2IPhHtoFsd-LdjMiQi1nn7kix4Nam0s,1011
 egglog/examples/fib.py,sha256=wAn-PjazxgHDkXAU4o2xTk_GtM_iGL0biV66vWM1st4,520
+egglog/examples/higher_order_functions.py,sha256=CBZqnqVdYGgztHw5QWtcnUhzs1nnBcw4O-gjSzAxjRc,1203
 egglog/examples/lambda_.py,sha256=hQBOaSw_yorNcbkQVu2EhgSc0IZNWIny7asaOlcUk9s,8496
 egglog/examples/matrix.py,sha256=_zmjgfFr2O_LjTcsTD-45_38Y_M1sP3AV39K6oFxAdw,5136
 egglog/examples/ndarrays.py,sha256=T-wwef-n-3LDSjaO35zA8AZH5DXFFqq0XBSCQKEXV6E,4186
@@ -18,7 +19,7 @@ egglog/examples/README.rst,sha256=QrbfmivODBvUvmY3-dHarcbC6bEvwoqAfTDhiI-aJxU,23
 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=diQNK-Uxxt6Qr58zJuQDM_BF6eKgwa-D5eYB5VJoT70,41412
+egglog/exp/array_api.py,sha256=b7MoUuH2sIW3KnUX7fGPaflwlCCskgrt1GsidVT-KG4,41474
 egglog/exp/array_api_jit.py,sha256=HIZzd0G17u-u_F4vfRdhoYvRo-ETx5HFO3RBcOfLcco,1287
 egglog/exp/array_api_numba.py,sha256=OjONBLdFRukx4vKiNK_rBxjMzsxbWpMEdD7JcGHJjmY,2924
 egglog/exp/array_api_program_gen.py,sha256=crgUYXXNhQdfTq31FSIpWLIfzNsgQD8ngg3OosCtIgg,19680
@@ -27,13 +28,13 @@ egglog/exp/siu_examples.py,sha256=KZmpkSCgbL4uqHhx2Jh9Adz1_cDbkIlyXPa1Kh_0O_8,78
 egglog/exp/__init__.py,sha256=G9zeKUcPBgIhgUg1meC86OfZVFETYIugyHWseTcCe_0,52
 egglog/graphviz_widget.py,sha256=YtI7LCFWHihDQ1qLvdj2SVYEcuZLSooFUYheunOTxdc,1339
 egglog/ipython_magic.py,sha256=VA19xAb6Sz7-IxlJBbnZW_gVFDqaYNnvdMB9QitndjE,1254
-egglog/pretty.py,sha256=4YzItMaXSQw2xzId9pnjUikOojwE6G7Hg0_M4sOGJ84,19071
+egglog/pretty.py,sha256=s5XBIB7kMrxV25aO-F4NW5Jli2DOoc72wwswFYxJPjM,19561
 egglog/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-egglog/runtime.py,sha256=P5wKThFglf2l5AynxgVAryyNjvLtFMbou7JCSSsORno,23013
-egglog/thunk.py,sha256=jYMb4owraT-3zGQEYkqMK76hk39Qdpc6bfEa7CVkjcg,1894
-egglog/type_constraint_solver.py,sha256=0c8oy-sLCVN0XndfAyRxO1AvIgfQuszw0NaRJ9ILysE,5692
+egglog/runtime.py,sha256=s32xdoD_mIdDRi_PJQ5E_kSNyhUkT310osH38kZfCSk,23174
+egglog/thunk.py,sha256=bm1zZUzwey0DijtMZkSulyfiWjCAg_1UmXaSG8qvvjs,2170
+egglog/type_constraint_solver.py,sha256=MW22gwryNOqsYaB__dCCjVvdnB1Km5fovoW2IgUXaVo,5836
 egglog/widget.css,sha256=WJS2M1wQdujhSTCakMa_ZXuoTPre1Uy1lPcvBE1LZQU,102
 egglog/widget.js,sha256=UNOER3sYZ-bS7Qhw9S6qtpR81FuHa5DzXQaWWtQq334,2021
 egglog/__init__.py,sha256=iUrVe5fb0XFyMCS3CwTjwhKEtU8KpIkdTpJpnUhm8o0,247
-egglog/bindings.cp310-win_amd64.pyd,sha256=pjYXlaSXtC50TWNdnGwLjKbntZ0WqfM5AR9hqehz2xc,4591104
-egglog-7.1.0.dist-info/RECORD,,
+egglog/bindings.cp310-win_amd64.pyd,sha256=NO6b-M4HPsi0uWxf6FKU7kBdS8WAsLbrJXeclR6XjuU,4591616
+egglog-7.2.0.dist-info/RECORD,,