egglog 9.0.0__cp313-cp313-win_amd64.whl

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

egglog/exp/__init__.py

@@ -0,0 +1,3 @@
+"""
+Experimental interfaces built on egglog.
+"""