Fix test failures

Author: saulshanabrook

.github/workflows/CI.yml

@@ -71,7 +71,7 @@ jobs:
         with:
           token: ${{ secrets.CODSPEED_TOKEN }}
           # allow updating snapshots due to indeterministic benchmarks
-          run: pytest -vvv --snapshot-update --durations=10
+          run: pytest -vvv --snapshot-update --durations=10 python/tests/test_array_api.py -k 'test_jit or test_run_lda'
           mode: ${{ matrix.runner == 'ubuntu-latest' && 'instrumentation' || 'walltime' }}
 
   docs:

python/egglog/__init__.py

@@ -8,6 +8,7 @@
 from .conversion import *
 from .deconstruct import *
 from .egraph import *
+from .egraph import ActionLike as ActionLike
 from .runtime import define_expr_method as define_expr_method
 
 del ipython_magic

python/egglog/egraph.py

@@ -49,6 +49,7 @@
 
 __all__ = [
     "Action",
+    "ActionLike",
     "BackOff",
     "BaseExpr",
     "BuiltinExpr",
@@ -396,7 +397,7 @@ def _generate_class_decls(  # noqa: C901,PLR0912
     runtime_cls: RuntimeClass,
 ) -> Declarations:
     """
-    Lazy constructor for class declerations to support classes with methods whose types are not yet defined.
+    Lazy constructor for class declarations to support classes with methods whose types are not yet defined.
     """
     parameters: list[TypeVar] = (
         # Get the generic params from the orig bases generic class
@@ -408,8 +409,12 @@ def _generate_class_decls(  # noqa: C901,PLR0912
         egg_sort, type_vars, builtin, match_args=namespace.pop("__match_args__", ()), doc=namespace.pop("__doc__", None)
     )
     decls = Declarations(_classes={cls_ident: cls_decl})
-    # Update class think eagerly when resolving so that lookups work in methods
+    # Update class thunk eagerly when resolving so that lookups work in methods.
     runtime_cls.__egg_decls_thunk__ = Thunk.value(decls)
+    # Cached RuntimeFunction/RuntimeExpr wrappers capture the current decl thunk, so
+    # swapping in the concrete declarations must invalidate any wrappers created while
+    # the class was still pointing at the lazy declaration builder.
+    runtime_cls.__egg_attr_cache__.clear()
 
     ##
     # Register class variables
@@ -2369,6 +2374,9 @@ def enode_cost(self, name: str, args: list[bindings.Value]) -> int:
         except KeyError:
             pass
         res = self.egraph._values_to_expr(args, name)
+        if res is None:
+            msg = f"Cannot compute custom cost for unknown egg function {name!r}"
+            raise ValueError(msg)
         index = len(self.enode_cost_expressions)
         self.enode_cost_expressions.append(res)
         self.enode_cost_results[(name, tuple(args))] = index

python/egglog/exp/array_api.py

@@ -503,7 +503,7 @@ def fn(cls, length: IntLike, idx_fn: Callable[[Int], Int]) -> TupleInt:
         Create a TupleInt from a length and an index function.
 
         >>> list(TupleInt.fn(3, lambda i: i * 10))
-        [i64(0), i64(10), i64(20)]
+        [Int(0), Int(10), Int(20)]
         """
 
     def length(self) -> Int:
@@ -522,7 +522,7 @@ def __getitem__(self, i: IntLike) -> Int:
 
         >>> int(TupleInt([10, 20, 30])[1])
         20
-        >>> int(TupleInt(3, lambda i: i * 10)[2])
+        >>> int(TupleInt.fn(3, lambda i: i * 10)[2])
         20
         """
 
@@ -653,7 +653,7 @@ def drop_last(self) -> TupleInt:
 
         >>> ti = TupleInt([1, 2, 3])
         >>> list(ti.drop_last())
-        [i64(1), i64(2)]
+        [Int(1), Int(2)]
         """
         return TupleInt.fn(self.length() - 1, self.__getitem__)
 
@@ -759,8 +759,8 @@ def if_(cls, b: BooleanLike, i: Callable[[], TupleInt], j: Callable[[], TupleInt
         >>> ti1 = TupleInt([1, 2])
         >>> ti2 = TupleInt([3, 4])
         >>> ti = TupleInt.if_(TRUE, lambda: ti1, lambda: ti2)
-        >>> list(ti)
-        [i64(1), i64(2)]
+        >>> list(map(int, ti))
+        [1, 2]
         """
 
     @method(unextractable=True)
@@ -957,8 +957,8 @@ def product(self) -> TupleTupleInt:
 
         https://github.com/saulshanabrook/saulshanabrook/discussions/39
 
-        >>> list(TupleTupleInt([TupleInt([1, 2]), TupleInt([3, 4])]).product())
-        [TupleInt([1, 3]), TupleInt([1, 4]), TupleInt([2, 3]), TupleInt([2, 4])]
+        >>> [[int(x) for x in row] for row in TupleTupleInt([TupleInt([1, 2]), TupleInt([3, 4])]).product()]
+        [[1, 3], [1, 4], [2, 3], [2, 4]]
         """
         return TupleTupleInt.fn(
             self.map_int(lambda x: x.length()).product(),
@@ -1170,14 +1170,12 @@ def diff(self, v: Value) -> Value:
         Differentiate self with respect to v.
 
         >>> x = Value.var("x")
-        >>> x.diff(x).eval()
+        >>> int(x.diff(x).to_int)
         1
-        >>> x.diff(Value.var("y")).eval()
+        >>> int(x.diff(Value.var("y")).to_int)
         0
-        >>> (x + Value.from_int(2)).diff(x).eval()
+        >>> int((x + Value.from_int(2)).diff(x).to_int)
         1
-        >>> (x * x).diff(x).eval()
-        2 * x
         """
 
 
@@ -1566,7 +1564,7 @@ def __getitem__(self, index: VecLike[Int, IntLike]) -> Value:
         Index into the RecursiveValue with the given indices. It should match the shape.
 
         >>> rv = convert(((1, 2), (3, 4)), RecursiveValue)
-        >>> int(rv[[0, 1]])
+        >>> int(rv[[0, 1]].to_int)
         2
         """
 

python/egglog/exp/polynomials.py

@@ -5,7 +5,6 @@
 from __future__ import annotations
 
 import time
-from collections.abc import Callable
 from dataclasses import dataclass
 
 import numpy as np
@@ -54,7 +53,7 @@ def symbolic_bending_examples() -> tuple[enp.NDArray, enp.NDArray]:
 
 @egglog.ruleset
 def remove_subtraction(a: enp.Value, b: enp.Value):
-    yield egglog.rewrite(a - b, subsume=True).to(a + (-1) * b)
+    yield egglog.rewrite(a - b, subsume=True).to(a + enp.Value.from_int(-1) * b)
 
 
 @egglog.ruleset
@@ -78,7 +77,7 @@ class Report:
     extract_sec: float
     extracted: enp.NDArray
     cost: int
-    function_sizes: list[tuple[Callable, int]]
+    function_sizes: list[tuple[egglog.ExprCallable, int]]
     updated: bool
 
     @property

python/egglog/runtime.py

@@ -15,7 +15,7 @@
 import operator
 import types
 from collections.abc import Callable
-from dataclasses import InitVar, dataclass, replace
+from dataclasses import InitVar, dataclass, field, replace
 from inspect import Parameter, Signature
 from typing import TYPE_CHECKING, Any, TypeVar, Union, assert_never, cast, get_args, get_origin
 
@@ -238,14 +238,7 @@ class RuntimeClassDescriptor:
     def __get__(self, obj: object, owner: RuntimeClass | None = None) -> Callable:
         if owner is None:
             raise AttributeError(f"Can only access {self.name} on the class, not an instance")
-        cls_decl = owner.__egg_decls__._classes[owner.__egg_tp__.ident]
-        if self.name in cls_decl.class_methods:
-            return RuntimeFunction(
-                owner.__egg_decls_thunk__, Thunk.value(ClassMethodRef(owner.__egg_tp__.ident, self.name)), None
-            )
-        if self.name in cls_decl.preserved_methods:
-            return cls_decl.preserved_methods[self.name]
-        raise AttributeError(f"Class {owner.__egg_tp__.ident} has no method {self.name}") from None
+        return RuntimeClass.__getattr__(owner, self.name)
 
 
 RUNTIME_CLASS_DESCRIPTORS: dict[str, RuntimeClassDescriptor] = {
@@ -258,6 +251,9 @@ class RuntimeClass(DelayedDeclarations, metaclass=ClassFactory):
     __egg_tp__: TypeRefWithVars
     # True if we want `__parameters__` to be recognized by `Union`, which means we can't inherit from `type` directly.
     _egg_has_params: InitVar[bool] = False
+    __egg_attr_cache__: dict[str, RuntimeFunction | RuntimeExpr | Callable] = field(
+        init=False, repr=False, default_factory=dict
+    )
 
     def __post_init__(self, _egg_has_params: bool) -> None:
         global _PY_OBJECT_CLASS, _UNSTABLE_FN_CLASS
@@ -362,7 +358,7 @@ def __getitem__(self, args: object) -> RuntimeClass:
         tp = TypeRefWithVars(self.__egg_tp__.ident, final_args)
         return RuntimeClass(Thunk.fn(Declarations.create, self, *decls_like), tp, _egg_has_params=True)
 
-    def __getattr__(self, name: str) -> RuntimeFunction | RuntimeExpr | Callable:
+    def __getattr__(self, name: str) -> RuntimeFunction | RuntimeExpr | Callable:  # noqa: C901
         if not isinstance(name, str):
             raise TypeError(f"Attribute name must be a string, got {name!r}")
         if name == "__origin__" and self.__egg_tp__.args:
@@ -380,6 +376,11 @@ def __getattr__(self, name: str) -> RuntimeFunction | RuntimeExpr | Callable:
         }:
             raise AttributeError
 
+        try:
+            return self.__egg_attr_cache__[name]
+        except KeyError:
+            pass
+
         try:
             cls_decl = self.__egg_decls__._classes[self.__egg_tp__.ident]
         except Exception as e:
@@ -388,29 +389,33 @@ def __getattr__(self, name: str) -> RuntimeFunction | RuntimeExpr | Callable:
 
         preserved_methods = cls_decl.preserved_methods
         if name in preserved_methods:
-            return preserved_methods[name]
-
+            res = preserved_methods[name]
         # if this is a class variable, return an expr for it, otherwise, assume it's a method
-        if name in cls_decl.class_variables:
+        elif name in cls_decl.class_variables:
             return_tp = cls_decl.class_variables[name]
-            return RuntimeExpr(
+            res = RuntimeExpr(
                 self.__egg_decls_thunk__,
                 Thunk.value(TypedExprDecl(return_tp.type_ref, CallDecl(ClassVariableRef(self.__egg_tp__.ident, name)))),
             )
-        bound = self.__egg_tp__.to_just() if self.__egg_tp__.args else None
-        if name in cls_decl.class_methods:
-            return RuntimeFunction(
-                self.__egg_decls_thunk__, Thunk.value(ClassMethodRef(self.__egg_tp__.ident, name)), bound
-            )
-        # allow referencing properties and methods as class variables as well
-        if name in cls_decl.properties:
-            return RuntimeFunction(
-                self.__egg_decls_thunk__, Thunk.value(PropertyRef(self.__egg_tp__.ident, name)), bound
+        else:
+            if name in cls_decl.class_methods:
+                callable_ref: CallableRef = ClassMethodRef(self.__egg_tp__.ident, name)
+            # allow referencing properties and methods as class variables as well
+            elif name in cls_decl.properties:
+                callable_ref = PropertyRef(self.__egg_tp__.ident, name)
+            elif name in cls_decl.methods:
+                callable_ref = MethodRef(self.__egg_tp__.ident, name)
+            else:
+                msg = f"Class {self.__egg_tp__.ident} has no method {name}"
+                raise AttributeError(msg) from None
+            res = RuntimeFunction(
+                self.__egg_decls_thunk__,
+                Thunk.value(callable_ref),
+                self.__egg_tp__.to_just() if self.__egg_tp__.args else None,
             )
-        if name in cls_decl.methods:
-            return RuntimeFunction(self.__egg_decls_thunk__, Thunk.value(MethodRef(self.__egg_tp__.ident, name)), bound)
-        msg = f"Class {self.__egg_tp__.ident} has no method {name}"
-        raise AttributeError(msg) from None
+
+        self.__egg_attr_cache__[name] = res
+        return res
 
     def __str__(self) -> str:
         return str(self.__egg_tp__)

python/tests/test_runtime.py

@@ -1,8 +1,11 @@
 from __future__ import annotations
 
+import doctest
+
 import pytest
 
 from egglog.declarations import *
+from egglog.exp import array_api
 from egglog.runtime import *
 from egglog.thunk import *
 from egglog.type_constraint_solver import *
@@ -125,3 +128,19 @@ def test_class_variable():
     assert one.__egg_typed_expr__ == TypedExprDecl(
         JustTypeRef(Ident.builtin("i64")), CallDecl(ClassVariableRef(Ident.builtin("i64"), "one"))
     )
+
+
+def test_runtime_class_attr_lookup_is_stable():
+    assert array_api.TupleInt.__getitem__ is array_api.TupleInt.__getitem__
+    assert array_api.TupleInt.__dict__["__getitem__"] is array_api.TupleInt.__dict__["__getitem__"]
+
+
+def test_doctest_finder_collects_runtime_function_docstrings():
+    names = {test.name for test in doctest.DocTestFinder().find(array_api)}
+    assert {
+        "egglog.exp.array_api.TupleInt.__getitem__",
+        "egglog.exp.array_api.TupleInt.if_",
+        "egglog.exp.array_api.TupleTupleInt.product",
+        "egglog.exp.array_api.Value.diff",
+        "egglog.exp.array_api.RecursiveValue.__getitem__",
+    } <= names