egglog 12.0.0__cp313-cp313t-manylinux_2_17_ppc64.manylinux2014_ppc64.whl

egglog/examples/matrix.py

@@ -0,0 +1,175 @@
+"""
+Matrix multiplication and Kronecker product.
+============================================
+"""
+
+from __future__ import annotations
+
+from egglog import *
+
+egraph = EGraph()
+
+
+class Dim(Expr):
+    """
+    A dimension of a matix.
+
+    >>> Dim(3) * Dim.named("n")
+    Dim(3) * Dim.named("n")
+    """
+
+    @method(egg_fn="Lit")
+    def __init__(self, value: i64Like) -> None: ...
+
+    @method(egg_fn="NamedDim")
+    @classmethod
+    def named(cls, name: StringLike) -> Dim:  # type: ignore[empty-body]
+        ...
+
+    @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),
+)
+
+
+class Matrix(Expr, egg_sort="MExpr"):
+    @method(egg_fn="Id")
+    @classmethod
+    def identity(cls, dim: Dim) -> Matrix:  # type: ignore[empty-body]
+        """
+        Create an identity matrix of the given dimension.
+        """
+
+    @method(egg_fn="NamedMat")
+    @classmethod
+    def named(cls, name: StringLike) -> Matrix:  # type: ignore[empty-body]
+        """
+        Create a named matrix.
+        """
+
+    @method(egg_fn="MMul")
+    def __matmul__(self, other: Matrix) -> Matrix:  # type: ignore[empty-body]
+        """
+        Matrix multiplication.
+        """
+
+    @method(egg_fn="nrows")
+    def nrows(self) -> Dim:  # type: ignore[empty-body]
+        """
+        Number of rows in the matrix.
+        """
+
+    @method(egg_fn="ncols")
+    def ncols(self) -> Dim:  # type: ignore[empty-body]
+        """
+        Number of columns in the matrix.
+        """
+
+
+@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(
+        A.ncols(),
+        A.nrows(),
+        B.ncols(),
+        B.nrows(),
+    ),
+    # demand rows and columns when we take the kronecker product
+    rule(eq(C).to(kron(A, B))).then(
+        A.ncols(),
+        A.nrows(),
+        B.ncols(),
+        B.nrows(),
+    ),
+)
+
+
+# Define a number of dimensions
+n = egraph.let("n", Dim.named("n"))
+m = egraph.let("m", Dim.named("m"))
+p = egraph.let("p", Dim.named("p"))
+
+# Define a number of matrices
+A = egraph.let("A", Matrix.named("A"))
+B = egraph.let("B", Matrix.named("B"))
+C = egraph.let("C", Matrix.named("C"))
+
+# Set each to be a square matrix of the given dimension
+egraph.register(
+    union(A.nrows()).with_(n),
+    union(A.ncols()).with_(n),
+    union(B.nrows()).with_(m),
+    union(B.ncols()).with_(m),
+    union(C.nrows()).with_(p),
+    union(C.ncols()).with_(p),
+)
+# Create an example which should equal the kronecker product of A and B
+ex1 = egraph.let("ex1", kron(Matrix.identity(n), B) @ kron(A, Matrix.identity(m)))
+rows = egraph.let("rows", ex1.nrows())
+cols = egraph.let("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
+simple_ex1 = egraph.let("simple_ex1", kron(A, B))
+egraph.check(eq(ex1).to(simple_ex1))
+
+ex2 = egraph.let("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

egglog/examples/multiset.py

@@ -0,0 +1,60 @@
+# mypy: disable-error-code="empty-body"
+"""
+Multiset example based off of egglog version
+============================================
+"""
+
+from __future__ import annotations
+
+from collections import Counter
+
+from egglog import *
+
+
+class Math(Expr):
+    def __init__(self, x: i64Like) -> None: ...
+
+
+@function
+def square(x: Math) -> Math: ...
+
+
+@ruleset
+def math_ruleset(i: i64):
+    yield rewrite(square(Math(i))).to(Math(i * i))
+
+
+egraph = EGraph()
+
+xs = MultiSet(Math(1), Math(2), Math(3))
+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).value) == Counter({Math(1): 1, Math(2): 1, Math(3): 1})
+
+
+inserted = MultiSet(Math(1), Math(2), Math(3), Math(4))
+egraph.register(inserted)
+egraph.check(xs.insert(Math(4)) == inserted)
+egraph.check(xs.contains(Math(1)))
+egraph.check(xs.not_contains(Math(4)))
+assert Math(1) in xs
+assert Math(4) not in xs
+
+egraph.check(xs.remove(Math(1)) == MultiSet(Math(2), Math(3)))
+
+assert egraph.extract(xs.length()).value == 3
+assert len(xs) == 3
+
+egraph.check(MultiSet(Math(1), Math(1)).length() == i64(2))
+
+egraph.check(MultiSet(Math(1)).pick() == Math(1))
+
+mapped = xs.map(square)
+egraph.register(mapped)
+egraph.run(math_ruleset)
+egraph.check(mapped == MultiSet(Math(1), Math(4), Math(9)))
+
+egraph.check(xs + xs == MultiSet(Math(1), Math(2), Math(3), Math(1), Math(2), Math(3)))

egglog/examples/ndarrays.py

@@ -0,0 +1,144 @@
+# mypy: disable-error-code="empty-body"
+
+"""
+N-Dimensional Arrays
+====================
+
+Example of building NDarray in the vein of Mathemetics of Arrays.
+"""
+
+from __future__ import annotations
+
+from egglog import *
+
+egraph = EGraph()
+
+
+class Value(Expr):
+    def __init__(self, v: i64Like) -> None: ...
+
+    def __mul__(self, other: Value) -> Value: ...
+
+    def __add__(self, other: Value) -> Value: ...
+
+
+i, j = vars_("i j", i64)
+egraph.register(
+    rewrite(Value(i) * Value(j)).to(Value(i * j)),
+    rewrite(Value(i) + Value(j)).to(Value(i + j)),
+)
+
+
+class Values(Expr):
+    def __init__(self, v: Vec[Value]) -> None: ...
+
+    def __getitem__(self, idx: Value) -> Value: ...
+
+    def length(self) -> Value: ...
+
+    def concat(self, other: Values) -> Values: ...
+
+
+@egraph.register
+def _values(vs: Vec[Value], other: Vec[Value]):
+    yield rewrite(Values(vs)[Value(i)]).to(vs[i])
+    yield rewrite(Values(vs).length()).to(Value(vs.length()))
+    yield rewrite(Values(vs).concat(Values(other))).to(Values(vs.append(other)))
+    # yield rewrite(l.concat(r).length()).to(l.length() + r.length())
+    # yield rewrite(l.concat(r)[idx])
+
+
+class NDArray(Expr):
+    """
+    An n-dimensional array.
+    """
+
+    def __getitem__(self, idx: Values) -> Value: ...
+
+    def shape(self) -> Values: ...
+
+
+@function
+def arange(n: Value) -> NDArray: ...
+
+
+@egraph.register
+def _ndarray_arange(n: Value, idx: Values):
+    yield rewrite(arange(n).shape()).to(Values(Vec(n)))
+    yield rewrite(arange(n)[idx]).to(idx[Value(0)])
+
+
+def assert_simplifies(left: Expr, right: Expr) -> None:
+    """
+    Simplify and print
+    """
+    egraph.register(left)
+    egraph.run(30)
+    res = egraph.extract(left)
+    print(f"{left} == {right} ➡  {res}")
+    egraph.check(eq(left).to(right))
+
+
+assert_simplifies(arange(Value(10)).shape(), Values(Vec(Value(10))))
+assert_simplifies(arange(Value(10))[Values(Vec(Value(0)))], Value(0))
+assert_simplifies(arange(Value(10))[Values(Vec(Value(1)))], Value(1))
+
+
+@function
+def py_value(s: StringLike) -> Value: ...
+
+
+@egraph.register
+def _py_value(l: String, r: String):
+    yield rewrite(py_value(l) + py_value(r)).to(py_value(join(l, " + ", r)))
+    yield rewrite(py_value(l) * py_value(r)).to(py_value(join(l, " * ", r)))
+
+
+@function
+def py_values(s: StringLike) -> Values: ...
+
+
+@egraph.register
+def _py_values(l: String, r: String):
+    yield rewrite(py_values(l)[py_value(r)]).to(py_value(join(l, "[", r, "]")))
+    yield rewrite(py_values(l).length()).to(py_value(join("len(", l, ")")))
+    yield rewrite(py_values(l).concat(py_values(r))).to(py_values(join(l, " + ", r)))
+
+
+@function
+def py_ndarray(s: StringLike) -> NDArray: ...
+
+
+@egraph.register
+def _py_ndarray(l: String, r: String):
+    yield rewrite(py_ndarray(l)[py_values(r)]).to(py_value(join(l, "[", r, "]")))
+    yield rewrite(py_ndarray(l).shape()).to(py_values(join(l, ".shape")))
+    yield rewrite(arange(py_value(l))).to(py_ndarray(join("np.arange(", l, ")")))
+
+
+assert_simplifies(py_ndarray("x").shape(), py_values("x.shape"))
+assert_simplifies(arange(py_value("x"))[py_values("y")], py_value("np.arange(x)[y]"))
+# assert_simplifies(arange(py_value("x"))[py_values("y")], py_value("y[0]"))
+
+
+@function
+def cross(l: NDArray, r: NDArray) -> NDArray: ...
+
+
+@egraph.register
+def _cross(l: NDArray, r: NDArray, idx: Values):
+    yield rewrite(cross(l, r).shape()).to(l.shape().concat(r.shape()))
+    yield rewrite(cross(l, r)[idx]).to(l[idx] * r[idx])
+
+
+assert_simplifies(cross(arange(Value(10)), arange(Value(11))).shape(), Values(Vec(Value(10), Value(11))))
+assert_simplifies(cross(py_ndarray("x"), py_ndarray("y")).shape(), py_values("x.shape + y.shape"))
+assert_simplifies(cross(py_ndarray("x"), py_ndarray("y"))[py_values("idx")], py_value("x[idx] * y[idx]"))
+
+
+@egraph.register
+def _cross_py(l: String, r: String):
+    yield rewrite(cross(py_ndarray(l), py_ndarray(r))).to(py_ndarray(join("np.multiply.outer(", l, ", ", r, ")")))
+
+
+assert_simplifies(cross(py_ndarray("x"), py_ndarray("y"))[py_values("idx")], py_value("np.multiply.outer(x, y)[idx]"))

egglog/examples/resolution.py

@@ -0,0 +1,84 @@
+"""
+Resolution theorem proving.
+===========================
+"""
+
+from __future__ import annotations
+
+from typing import ClassVar
+
+from egglog import *
+
+egraph = EGraph()
+
+
+class Boolean(Expr):
+    FALSE: ClassVar[Boolean]
+
+    def __or__(self, other: Boolean) -> Boolean:  # type: ignore[empty-body]
+        ...
+
+    def __invert__(self) -> Boolean:  # type: ignore[empty-body]
+        ...
+
+
+# Show off two ways of creating constants, either as top level values or as classvars
+T = constant("T", Boolean)
+F = Boolean.FALSE
+
+p, a, b, c, as_, bs = vars_("p a b c as bs", Boolean)
+egraph.register(
+    # clauses are assumed in the normal form (or a (or b (or c False)))
+    union(~F).with_(T),
+    union(~T).with_(F),
+    # "Solving" negation equations
+    rule(eq(~p).to(T)).then(union(p).with_(F)),
+    rule(eq(~p).to(F)).then(union(p).with_(T)),
+    # canonicalize associtivity. "append" for clauses terminate with false
+    rewrite((a | b) | c).to(a | (b | c)),
+    # commutativity
+    rewrite(a | (b | c)).to(b | (a | c)),
+    # absoprtion
+    rewrite(a | (a | b)).to(a | b),
+    rewrite(a | (~a | b)).to(T),
+    # Simplification
+    rewrite(F | a).to(a),
+    rewrite(a | F).to(a),
+    rewrite(T | a).to(T),
+    rewrite(a | T).to(T),
+    # unit propagation
+    # This is kind of interesting actually.
+    # Looks a bit like equation solving
+    rule(eq(T).to(p | F)).then(union(p).with_(T)),
+    # resolution
+    # This counts on commutativity to bubble everything possible up to the front of the clause.
+    rule(
+        eq(T).to(a | as_),
+        eq(T).to(~a | bs),
+    ).then(
+        union(as_ | bs).with_(T),
+    ),
+)
+
+
+# Example predicate
+@function
+def pred(x: i64Like) -> Boolean:  # type: ignore[empty-body]
+    ...
+
+
+p0 = egraph.let("p0", pred(0))
+p1 = egraph.let("p1", pred(1))
+p2 = egraph.let("p2", pred(2))
+egraph.register(
+    union(p1 | (~p2 | F)).with_(T),
+    union(p2 | (~p0 | F)).with_(T),
+    union(p0 | (~p1 | F)).with_(T),
+    union(p1).with_(F),
+    union(~p0 | (~p1 | (p2 | F))).with_(T),
+)
+egraph.run(10)
+egraph.check(ne(T).to(F))
+egraph.check(eq(p0).to(F))
+egraph.check(eq(p2).to(F))
+egraph

egglog/examples/schedule_demo.py

@@ -0,0 +1,34 @@
+"""
+Schedule demo
+=============
+"""
+
+from __future__ import annotations
+
+from egglog import *
+
+left = relation("left", i64)
+right = relation("right", i64)
+
+
+x, y = vars_("x y", i64)
+step_left = ruleset(
+    rule(
+        left(x),
+        right(x),
+    ).then(left(x + 1)),
+)
+step_right = ruleset(
+    rule(
+        left(x),
+        right(y),
+        eq(x).to(y + 1),
+    ).then(right(x)),
+)
+
+egraph = EGraph()
+egraph.register(left(i64(0)), right(i64(0)))
+egraph.run((step_right.saturate() + step_left.saturate()) * 10)
+egraph.check(left(i64(10)), right(i64(9)))
+egraph.check_fail(left(i64(11)), right(i64(10)))
+egraph