egglog 0.4.0__cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

egglog/examples/README.rst

@@ -0,0 +1,5 @@
+Examples Gallery
+================
+
+This is a gallery of examples, most of which were translated from the original
+`egglog rust examples <https://github.com/egraphs-good/egglog/tree/08a6e8fecdb77e6ba72a1b1d9ff4aff33229912c/tests>`_.

egglog/examples/eqsat_basic.py

@@ -0,0 +1,43 @@
+"""
+Basic equality saturation example.
+==================================
+"""
+from __future__ import annotations
+
+from egglog import *
+
+egraph = EGraph()
+
+
+@egraph.class_
+class Num(BaseExpr):
+    def __init__(self, value: i64Like) -> None:
+        ...
+
+    @classmethod
+    def var(cls, name: StringLike) -> Num:  # type: ignore[empty-body]
+        ...
+
+    def __add__(self, other: Num) -> Num:  # type: ignore[empty-body]
+        ...
+
+    def __mul__(self, other: Num) -> Num:  # type: ignore[empty-body]
+        ...
+
+
+# expr1 = 2 * (x + 3)
+expr1 = egraph.define("expr1", Num(2) * (Num.var("x") + Num(3)))
+# expr2 = 6 + 2 * x
+expr2 = egraph.define("expr2", Num(6) + Num(2) * Num.var("x"))
+
+a, b, c = vars_("a b c", Num)
+i, j = vars_("i j", i64)
+egraph.register(
+    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)),
+)
+egraph.run(10)
+egraph.check(eq(expr1).to(expr2))
+egraph

egglog/examples/fib.py

@@ -0,0 +1,28 @@
+"""
+Fibonacci numbers example
+=========================
+"""
+from __future__ import annotations
+
+from egglog import *
+
+egraph = EGraph()
+
+
+@egraph.function
+def fib(x: i64Like) -> i64:  # type: ignore[empty-body]
+    ...
+
+
+f0, f1, x = vars_("f0 f1 x", i64)
+egraph.register(
+    set_(fib(0)).to(i64(1)),
+    set_(fib(1)).to(i64(1)),
+    rule(
+        eq(f0).to(fib(x)),
+        eq(f1).to(fib(x + 1)),
+    ).then(set_(fib(x + 2)).to(f0 + f1)),
+)
+egraph.run(7)
+egraph.check(eq(fib(i64(7))).to(i64(21)))
+egraph

egglog/examples/lambda.py

@@ -0,0 +1,310 @@
+"""
+Lambda Calculus
+===============
+"""
+# mypy: disable-error-code=empty-body
+from __future__ import annotations
+
+from typing import Callable, ClassVar
+
+from egglog import *
+
+egraph = EGraph()
+
+# TODO: Debug extracting constants
+
+
+@egraph.class_
+class Val(BaseExpr):
+    """
+    A value is a number or a boolean.
+    """
+
+    TRUE: ClassVar[Val]
+    FALSE: ClassVar[Val]
+
+    def __init__(self, v: i64Like) -> None:
+        ...
+
+
+@egraph.class_
+class Var(BaseExpr):
+    def __init__(self, v: StringLike) -> None:
+        ...
+
+
+@egraph.class_
+class Term(BaseExpr):
+    @classmethod
+    def val(cls, v: Val) -> Term:
+        ...
+
+    @classmethod
+    def var(cls, v: Var) -> Term:
+        ...
+
+    def __add__(self, other: Term) -> Term:
+        ...
+
+    def __eq__(self, other: Term) -> Term:  # type: ignore[override]
+        ...
+
+    def __call__(self, other: Term) -> Term:
+        ...
+
+    def eval(self) -> Val:
+        ...
+
+    def v(self) -> Var:
+        ...
+
+
+@egraph.function
+def lam(x: Var, t: Term) -> Term:
+    ...
+
+
+@egraph.function
+def let_(x: Var, t: Term, b: Term) -> Term:
+    ...
+
+
+@egraph.function
+def fix(x: Var, t: Term) -> Term:
+    ...
+
+
+@egraph.function
+def if_(c: Term, t: Term, f: Term) -> Term:
+    ...
+
+
+StringSet = Set[Var]
+
+
+@egraph.function(merge=lambda old, new: old & new)
+def freer(t: Term) -> StringSet:
+    ...
+
+
+(v, v1, v2) = vars_("v v1 v2", Val)
+(t, t1, t2, t3, t4) = vars_("t t1 t2 t3 t4", Term)
+(x, y) = vars_("x y", Var)
+fv, fv1, fv2, fv3 = vars_("fv fv1 fv2 fv3", StringSet)
+i1, i2 = vars_("i1 i2", i64)
+egraph.register(
+    # freer
+    rule(eq(t).to(Term.val(v))).then(set_(freer(t)).to(StringSet.empty())),
+    rule(eq(t).to(Term.var(x))).then(set_(freer(t)).to(StringSet.empty().insert(x))),
+    rule(eq(t).to(t1 + t2), eq(freer(t1)).to(fv1), eq(freer(t2)).to(fv2)).then(set_(freer(t)).to(fv1 | fv2)),
+    rule(eq(t).to(t1 == t2), eq(freer(t1)).to(fv1), eq(freer(t2)).to(fv2)).then(set_(freer(t)).to(fv1 | fv2)),
+    rule(eq(t).to(t1(t2)), eq(freer(t1)).to(fv1), eq(freer(t2)).to(fv2)).then(set_(freer(t)).to(fv1 | fv2)),
+    rule(eq(t).to(lam(x, t1)), eq(freer(t1)).to(fv)).then(set_(freer(t)).to(fv.remove(x))),
+    rule(eq(t).to(let_(x, t1, t2)), eq(freer(t1)).to(fv1), eq(freer(t2)).to(fv2)).then(
+        set_(freer(t)).to(fv1.remove(x) | fv2)
+    ),
+    rule(eq(t).to(fix(x, t1)), eq(freer(t1)).to(fv)).then(set_(freer(t)).to(fv.remove(x))),
+    rule(eq(t).to(if_(t1, t2, t3)), eq(freer(t1)).to(fv1), eq(freer(t2)).to(fv2), eq(freer(t3)).to(fv3)).then(
+        set_(freer(t)).to(fv1 | fv2 | fv3)
+    ),
+    # eval
+    rule(eq(t).to(Term.val(v))).then(set_(t.eval()).to(v)),
+    rule(eq(t).to(t1 + t2), eq(Val(i1)).to(t1.eval()), eq(Val(i2)).to(t2.eval())).then(
+        union(t.eval()).with_(Val(i1 + i2))
+    ),
+    rule(eq(t).to(t1 == t2), eq(t1.eval()).to(t2.eval())).then(union(t.eval()).with_(Val.TRUE)),
+    rule(eq(t).to(t1 == t2), eq(t1.eval()).to(v1), eq(t2.eval()).to(v2), v1 != v2).then(
+        union(t.eval()).with_(Val.FALSE)
+    ),
+    rule(eq(v).to(t.eval())).then(union(t).with_(Term.val(v))),
+    # if
+    rewrite(if_(Term.val(Val.TRUE), t1, t2)).to(t1),
+    rewrite(if_(Term.val(Val.FALSE), t1, t2)).to(t2),
+    # if-elim
+    # Adds let rules so next one can match on them
+    rule(eq(t).to(if_(Term.var(x) == t1, t2, t3))).then(let_(x, t1, t2), let_(x, t1, t3)),
+    rewrite(if_(Term.var(x) == t1, t2, t3)).to(
+        t3,
+        eq(let_(x, t1, t2)).to(let_(x, t1, t3)),
+    ),
+    # add-comm
+    rewrite(t1 + t2).to(t2 + t1),
+    # add-assoc
+    rewrite((t1 + t2) + t3).to(t1 + (t2 + t3)),
+    # eq-comm
+    rewrite(t1 == t2).to(t2 == t1),
+    # Fix
+    rewrite(fix(x, t)).to(let_(x, fix(x, t), t)),
+    # beta reduction
+    rewrite(lam(x, t)(t1)).to(let_(x, t1, t)),
+    # let-app
+    rewrite(let_(x, t, t1(t2))).to(let_(x, t, t1)(let_(x, t, t2))),
+    # let-add
+    rewrite(let_(x, t, t1 + t2)).to(let_(x, t, t1) + let_(x, t, t2)),
+    # let-eq
+    rewrite(let_(x, t, t1 == t2)).to(let_(x, t, t1) == let_(x, t, t2)),
+    # let-const
+    rewrite(let_(x, t, Term.val(v))).to(Term.val(v)),
+    # let-if
+    rewrite(let_(x, t, if_(t1, t2, t3))).to(if_(let_(x, t, t1), let_(x, t, t2), let_(x, t, t3))),
+    # let-var-same
+    rewrite(let_(x, t, Term.var(x))).to(t),
+    # let-var-diff
+    rewrite(let_(x, t, Term.var(y))).to(Term.var(y), x != y),
+    # let-lam-same
+    rewrite(let_(x, t, lam(x, t1))).to(lam(x, t1)),
+    # let-lam-diff
+    rewrite(let_(x, t, lam(y, t1))).to(lam(y, let_(x, t, t1)), x != y, eq(fv).to(freer(t)), fv.not_contains(y)),
+    rule(eq(t).to(let_(x, t1, lam(y, t2))), x != y, eq(fv).to(freer(t1)), fv.contains(y)).then(
+        union(t).with_(lam(t.v(), let_(x, t1, let_(y, Term.var(t.v()), t2))))
+    ),
+)
+
+result = egraph.relation("result")
+
+
+def l(fn: Callable[[Term], Term]) -> Term:  # noqa
+    """
+    Create a lambda term from a function
+    """
+    # Use first var name from fn
+    x = fn.__code__.co_varnames[0]
+    return lam(Var(x), fn(Term.var(Var(x))))
+
+
+def assert_simplifies(left: BaseExpr, right: BaseExpr) -> None:
+    """
+    Simplify and print
+    """
+    with egraph:
+        egraph.register(left)
+        egraph.run(30)
+        res = egraph.extract(left)
+        print(f"{left} ➡  {res}")
+        egraph.check(eq(left).to(right))
+
+
+assert_simplifies((Term.val(Val(1))).eval(), Val(1))
+assert_simplifies((Term.val(Val(1)) + Term.val(Val(2))).eval(), Val(3))
+
+
+# lambda under
+assert_simplifies(
+    l(lambda x: Term.val(Val(4)) + l(lambda y: y)(Term.val(Val(4)))),
+    l(lambda x: Term.val(Val(8))),
+)
+# lambda if elim
+a = Term.var(Var("a"))
+b = Term.var(Var("b"))
+with egraph:
+    e1 = egraph.define("e1", if_(a == b, a + a, a + b))
+    egraph.run(10)
+    egraph.check(eq(e1).to(a + b))
+
+# lambda let simple
+x = Var("x")
+y = Var("y")
+assert_simplifies(
+    let_(x, Term.val(Val(0)), let_(y, Term.val(Val(1)), Term.var(x) + Term.var(y))),
+    Term.val(Val(1)),
+)
+# lambda capture
+assert_simplifies(
+    let_(x, Term.val(Val(1)), l(lambda x: x)),
+    l(lambda x: x),
+)
+# lambda capture free
+with egraph:
+    e5 = egraph.define("e5", let_(y, Term.var(x) + Term.var(x), l(lambda x: Term.var(y))))
+    egraph.run(10)
+    egraph.check(freer(l(lambda x: Term.var(y))).contains(y))
+    egraph.check_fail(eq(e5).to(l(lambda x: x + x)))
+
+# lambda_closure_not_seven
+with egraph:
+    e6 = egraph.define(
+        "e6",
+        let_(
+            Var("five"),
+            Term.val(Val(5)),
+            let_(
+                Var("add-five"),
+                l(lambda x: x + Term.var(Var("five"))),
+                let_(Var("five"), Term.val(Val(6)), Term.var(Var("add-five"))(Term.val(Val(1)))),
+            ),
+        ),
+    )
+    egraph.run(10)
+    egraph.check_fail(eq(e6).to(Term.val(Val(7))))
+    egraph.check(eq(e6).to(Term.val(Val(6))))
+
+
+# lambda_compose
+with egraph:
+    compose = Var("compose")
+    add1 = Var("add1")
+    e7 = egraph.define(
+        "e7",
+        let_(
+            compose,
+            l(
+                lambda f: l(
+                    lambda g: l(
+                        lambda x: f(g(x)),
+                    ),
+                ),
+            ),
+            let_(
+                add1,
+                l(lambda y: y + Term.val(Val(1))),
+                Term.var(compose)(Term.var(add1))(Term.var(add1)),
+            ),
+        ),
+    )
+    egraph.run(20)
+    egraph.register(
+        rule(
+            eq(t1).to(l(lambda x: Term.val(Val(1)) + l(lambda y: Term.val(Val(1)) + y)(x))),
+            eq(t2).to(l(lambda x: x + Term.val(Val(2)))),
+        ).then(result())
+    )
+    egraph.run(1)
+    egraph.check(result())
+
+
+# lambda_if_simple
+assert_simplifies(if_(Term.val(Val(1)) == Term.val(Val(1)), Term.val(Val(7)), Term.val(Val(9))), Term.val(Val(7)))
+
+
+# # lambda_compose_many
+assert_simplifies(
+    let_(
+        compose,
+        l(lambda f: l(lambda g: l(lambda x: f(g(x))))),
+        let_(
+            add1,
+            l(lambda y: y + Term.val(Val(1))),
+            Term.var(compose)(Term.var(add1))(
+                Term.var(compose)(Term.var(add1))(
+                    Term.var(compose)(Term.var(add1))(
+                        Term.var(compose)(Term.var(add1))(
+                            Term.var(compose)(Term.var(add1))(Term.var(compose)(Term.var(add1))(Term.var(add1)))
+                        )
+                    )
+                )
+            ),
+        ),
+    ),
+    l(lambda x: x + Term.val(Val(7))),
+)
+
+# lambda_if
+zeroone = Var("zeroone")
+assert_simplifies(
+    let_(
+        zeroone,
+        l(lambda x: if_(x == Term.val(Val(0)), Term.val(Val(0)), Term.val(Val(1)))),
+        Term.var(zeroone)(Term.val(Val(0))) + Term.var(zeroone)(Term.val(Val(10))),
+    ),
+    Term.val(Val(1)),
+)

egglog/examples/matrix.py

@@ -0,0 +1,184 @@
+"""
+Matrix multiplication and Kronecker product.
+============================================
+"""
+from __future__ import annotations
+
+from egglog import *
+
+egraph = EGraph()
+
+
+@egraph.class_
+class Dim(BaseExpr):
+    """
+    A dimension of a matix.
+
+    >>> Dim(3) * Dim.named("n")
+    Dim(3) * Dim.named("n")
+    """
+
+    @egraph.method(egg_fn="Lit")
+    def __init__(self, value: i64Like) -> None:
+        ...
+
+    @egraph.method(egg_fn="NamedDim")
+    @classmethod
+    def named(cls, name: StringLike) -> Dim:  # type: ignore[empty-body]
+        ...
+
+    @egraph.method(egg_fn="Times")
+    def __mul__(self, other: Dim) -> Dim:  # type: ignore[empty-body]
+        ...
+
+
+a, b, c, n = vars_("a b c n", Dim)
+i, j = vars_("i j", i64)
+egraph.register(
+    rewrite(a * (b * c)).to((a * b) * c),
+    rewrite((a * b) * c).to(a * (b * c)),
+    rewrite(Dim(i) * Dim(j)).to(Dim(i * j)),
+    rewrite(a * b).to(b * a),
+)
+
+
+@egraph.class_(egg_sort="MExpr")
+class Matrix(BaseExpr):
+    @egraph.method(egg_fn="Id")
+    @classmethod
+    def identity(cls, dim: Dim) -> Matrix:  # type: ignore[empty-body]
+        """
+        Create an identity matrix of the given dimension.
+        """
+        ...
+
+    @egraph.method(egg_fn="NamedMat")
+    @classmethod
+    def named(cls, name: StringLike) -> Matrix:  # type: ignore[empty-body]
+        """
+        Create a named matrix.
+        """
+        ...
+
+    @egraph.method(egg_fn="MMul")
+    def __matmul__(self, other: Matrix) -> Matrix:  # type: ignore[empty-body]
+        """
+        Matrix multiplication.
+        """
+        ...
+
+    @egraph.method(egg_fn="nrows")
+    def nrows(self) -> Dim:  # type: ignore[empty-body]
+        """
+        Number of rows in the matrix.
+        """
+        ...
+
+    @egraph.method(egg_fn="ncols")
+    def ncols(self) -> Dim:  # type: ignore[empty-body]
+        """
+        Number of columns in the matrix.
+        """
+        ...
+
+
+@egraph.function(egg_fn="Kron")
+def kron(a: Matrix, b: Matrix) -> Matrix:  # type: ignore[empty-body]
+    """
+    Kronecker product of two matrices.
+
+    https://en.wikipedia.org/wiki/Kronecker_product#Definition
+    """
+    ...
+
+
+A, B, C, D = vars_("A B C D", Matrix)
+egraph.register(
+    # The dimensions of a kronecker product are the product of the dimensions
+    rewrite(kron(A, B).nrows()).to(A.nrows() * B.nrows()),
+    rewrite(kron(A, B).ncols()).to(A.ncols() * B.ncols()),
+    # The dimensions of a matrix multiplication are the number of rows of the first
+    # matrix and the number of columns of the second matrix.
+    rewrite((A @ B).nrows()).to(A.nrows()),
+    rewrite((A @ B).ncols()).to(B.ncols()),
+    # The dimensions of an identity matrix are the input dimension
+    rewrite(Matrix.identity(n).nrows()).to(n),
+    rewrite(Matrix.identity(n).ncols()).to(n),
+)
+egraph.register(
+    # Multiplication by an identity matrix is the same as the other matrix
+    rewrite(Matrix.identity(n) @ A).to(A),
+    rewrite(A @ Matrix.identity(n)).to(A),
+    # Matrix multiplication is associative
+    rewrite(A @ (B @ C)).to((A @ B) @ C),
+    rewrite((A @ B) @ C).to(A @ (B @ C)),
+    # Kronecker product is associative
+    rewrite(kron(A, kron(B, C))).to(kron(kron(A, B), C)),
+    rewrite(kron(kron(A, B), C)).to(kron(A, kron(B, C))),
+    # Kronecker product distributes over matrix multiplication
+    rewrite(kron(A @ C, B @ D)).to(kron(A, B) @ kron(C, D)),
+    rewrite(kron(A, B) @ kron(C, D)).to(
+        kron(A @ C, B @ D),
+        # Only when the dimensions match
+        eq(A.ncols()).to(C.nrows()),
+        eq(B.ncols()).to(D.nrows()),
+    ),
+)
+egraph.register(
+    # demand rows and columns when we multiply matrices
+    rule(eq(C).to(A @ B)).then(
+        let("demand1", A.ncols()),
+        let("demand2", A.nrows()),
+        let("demand3", B.ncols()),
+        let("demand4", B.nrows()),
+    ),
+    # demand rows and columns when we take the kronecker product
+    rule(eq(C).to(kron(A, B))).then(
+        let("demand1", A.ncols()),
+        let("demand2", A.nrows()),
+        let("demand3", B.ncols()),
+        let("demand4", B.nrows()),
+    ),
+)
+
+
+# Define a number of dimensions
+n = egraph.define("n", Dim.named("n"))
+m = egraph.define("m", Dim.named("m"))
+p = egraph.define("p", Dim.named("p"))
+
+# Define a number of matrices
+A = egraph.define("A", Matrix.named("A"))
+B = egraph.define("B", Matrix.named("B"))
+C = egraph.define("C", Matrix.named("C"))
+
+# Set each to be a square matrix of the given dimension
+egraph.register(
+    set_(A.nrows()).to(n),
+    set_(A.ncols()).to(n),
+    set_(B.nrows()).to(m),
+    set_(B.ncols()).to(m),
+    set_(C.nrows()).to(p),
+    set_(C.ncols()).to(p),
+)
+# Create an example which should equal the kronecker product of A and B
+ex1 = egraph.define("ex1", kron(Matrix.identity(n), B) @ kron(A, Matrix.identity(m)))
+rows = egraph.define("rows", ex1.nrows())
+cols = egraph.define("cols", ex1.ncols())
+
+egraph.run(20)
+
+egraph.check(eq(B.nrows()).to(m))
+egraph.check(eq(kron(Matrix.identity(n), B).nrows()).to(n * m))
+
+# Verify it matches the expected result
+# TODO
+simple_ex1 = egraph.define("simple_ex1", kron(A, B))
+egraph.check(eq(ex1).to(simple_ex1))
+
+ex2 = egraph.define("ex2", kron(Matrix.identity(p), C) @ kron(A, Matrix.identity(m)))
+
+egraph.run(10)
+# Verify it is not simplified
+egraph.check_fail(eq(ex2).to(kron(A, C)))
+egraph