Add ndarray demo

Author: saulshanabrook

Cargo.lock

@@ -12,7 +12,7 @@ crate-type = ["cdylib"]
 # Use unreleased version to depend on signature improvements https://github.com/PyO3/pyo3/pull/2702
 pyo3 = { version = "0.18.1", features = ["extension-module"] }
 # egglog = { git = "https://github.com/egraphs-good/egglog", rev = "39b199d9bfce9cc47d0c54977279c5b04231e717" }
-egglog = { git = "https://github.com/saulshanabrook/egg-smol", rev = "de202c4930a9a983361bd7af10f8d6c1a3740c08" }
+egglog = { git = "https://github.com/saulshanabrook/egg-smol", rev = "989f0527d46550f05562a3a8c0e87a0b1280a930" }
 
 # egglog = { path = "../egg-smol" }
 pyo3-log = "0.8.1"

Cargo.toml

@@ -1,16 +1,5 @@
 {
  "cells": [
-  {
-   "cell_type": "markdown",
-   "id": "a74307d3-b810-4f9d-9663-d5b0030357ae",
-   "metadata": {},
-   "source": [
-    "* Campy\n",
-    "* Add more complicated example\n",
-    "* Change fonts\n",
-    "* Make sure examples are done"
-   ]
-  },
   {
    "cell_type": "markdown",
    "id": "a611c138-1afd-4d6f-9578-4f358d2438eb",

Presentation.ipynb

@@ -11,9 +11,15 @@ _This project uses semantic versioning. Before 1.0.0, this means that every brea
   - Added `Set` sort and removed set method from `Map`
   - Added `Vec` sort
   - Added support for variable args for builtin functions, to use in creation of `Vec` and `Set` sorts.
+  - Added suport for joining `String`s
 - Switch generated egg names to use `.` as seperate (i.e. `Math.__add__`) instead of `_` (i.e. `Math___add__`)
 - Adds support for modules to define functions/sorts/rules without executing them, for reuse in other modules
   - Moved simplifying and running rulesets to the `run` and `simplify` methods on the `EGraph` from those methods on the `Ruleset` since we can now create `Rulset`s for modules which don't have an EGraph attached and can't be run
+- Fixed extracting classmethods which required generic args to cls
+- Added support for alternative way of creating variables using functions
+- Add NDarray example
+- Render EGraphs with `graphviz` in the notebook
+- Add `%%egglog` magic to the notebook
 
 ## 0.4.0 (2023-05-03)
 

docs/changelog.md

@@ -328,6 +328,21 @@ Since it uses a fluent API, static type checkers can verify that the type of the
 
 The `(birewrite ...)` command in egglog is syntactic sugar for creating two rewrites, one in each direction. In Python, we can use the `birewrite(expr).to(expr, *when)` function to create two rules that rewrite in each direction.
 
+### Using funcitons to define vars
+
+Instead of defining variables with `vars_`, we can also use functions to define variables. This can be more succinct
+and also will make sure the variables won't be used outside of the scope of the function.
+
+```{code-cell} python
+# egg: (rewrite (Mul a b) (Mul b a))
+# egg: (rewrite (Add a b) (Add b a))
+
+@egraph.register
+def _math(a: Math, b: Math)
+    yield rewrite(a * b).to(b * a)
+    yield rewrite(a + b).to(b + a)
+```
+
 ## Running
 
 To run the egraph, we can use the `egraph.run()` function. This will run until a fixed point is reached, or until a timeout is reached.

docs/reference/egglog-translation.md

@@ -21,6 +21,7 @@
     "Rational",
     "Set",
     "Vec",
+    "join",
 ]
 
 
@@ -147,13 +148,19 @@ def max(self, other: f64Like) -> f64:  # type: ignore[empty-body]
         ...
 
 
+StringLike = Union[str, "String"]
+
+
 @BUILTINS.class_
 class String(BaseExpr):
     def __init__(self, value: str):
         ...
 
 
-StringLike = Union[str, String]
+@BUILTINS.function(egg_fn="+")
+def join(*strings: StringLike) -> String:  # type: ignore[empty-body]
+    ...
+
 
 T = TypeVar("T", bound=BaseExpr)
 V = TypeVar("V", bound=BaseExpr)
@@ -328,3 +335,11 @@ def not_contains(self, value: T) -> Unit:  # type: ignore[empty-body]
     @BUILTINS.method(egg_fn="vec-contains")
     def contains(self, value: T) -> Unit:  # type: ignore[empty-body]
         ...
+
+    @BUILTINS.method(egg_fn="vec-length")
+    def length(self) -> i64:  # type: ignore[empty-body]
+        ...
+
+    @BUILTINS.method(egg_fn="vec-get")
+    def __getitem__(self, index: i64Like) -> T:  # type: ignore[empty-body]
+        ...

python/egglog/builtins.py

@@ -536,7 +536,18 @@ def from_egg(cls, mod_decls: ModuleDeclarations, call: bindings.Call) -> TypedEx
 
         # Find the first callable ref that matches the call
         for callable_ref in mod_decls.get_callable_refs(call.name):
-            tcs = TypeConstraintSolver()
+            # If this is a classmethod, we might need the type params that were bound for this type
+            # egglog currently only allows one instantiated type of any generic sort to be used in any program
+            # So we just lookup what args were registered for thsi sort
+            if isinstance(callable_ref, ClassMethodRef):
+                for registered_tp in mod_decls._decl._type_ref_to_egg_sort.keys():
+                    if registered_tp.name == callable_ref.class_name:
+                        tcs = TypeConstraintSolver.from_type_parameters(registered_tp.args)
+                        break
+                else:
+                    raise ValueError(f"Could not find type parameters for class {callable_ref.class_name}")
+            else:
+                tcs = TypeConstraintSolver()
             fn_decl = mod_decls.get_function_decl(callable_ref)
             return_tp = tcs.infer_return_type(fn_decl.arg_types, fn_decl.return_type, fn_decl.var_arg_type, arg_types)
             return TypedExprDecl(return_tp, cls(callable_ref, tuple(results)))

python/egglog/declarations.py

@@ -23,6 +23,7 @@
     overload,
 )
 
+import graphviz
 from egglog.declarations import Declarations
 from typing_extensions import ParamSpec, get_args, get_origin
 
@@ -214,7 +215,10 @@ def _class(
                 # If this is an i64, use the runtime class for the alias so that i64Like is resolved properly
                 # Otherwise, this might be a Map in which case pass in the original cls so that we
                 # can do Map[T, V] on it, which is not allowed on the runtime class
-                cls_type_and_name=(RuntimeClass(self._mod_decls, "i64") if cls_name == "i64" else cls, cls_name),
+                cls_type_and_name=(
+                    RuntimeClass(self._mod_decls, cls_name) if cls_name in {"i64", "String"} else cls,
+                    cls_name,
+                ),
             )
 
         # Register != as a method so we can print it as a string
@@ -450,10 +454,15 @@ def _resolve_type_annotation(
             return class_to_ref(tp).to_var()
         raise TypeError(f"Unexpected type annotation {tp}")
 
-    def register(self, *commands: CommandLike) -> None:
+    def register(self, command_or_generator: CommandLike | CommandGenerator, *commands: CommandLike) -> None:
         """
         Registers any number of rewrites or rules.
         """
+        if isinstance(command_or_generator, FunctionType):
+            assert not commands
+            commands = tuple(_command_generator(command_or_generator))
+        else:
+            commands = (cast(CommandLike, command_or_generator), *commands)
         self._process_commands(_command_like(command)._to_egg_command(self._mod_decls) for command in commands)
 
     def ruleset(self, name: str) -> Ruleset:
@@ -587,9 +596,12 @@ def _repr_mimebundle_(self, *args, **kwargs):
         """
         Returns the graphviz representation of the e-graph.
         """
-        import graphviz
 
-        return graphviz.Source(self._egraph.to_graphviz_string())._repr_mimebundle_(*args, **kwargs)
+        return self.graphviz._repr_mimebundle_(*args, **kwargs)
+
+    @property
+    def graphviz(self) -> graphviz.Source:
+        return graphviz.Source(self._egraph.to_graphviz_string())
 
     def display(self):
         """
@@ -982,6 +994,18 @@ def _command_like(command_like: CommandLike) -> Command:
     return command_like
 
 
+CommandGenerator = Callable[..., Iterable[Command]]
+
+
+def _command_generator(gen: CommandGenerator) -> Iterable[Command]:
+    """
+    Calls the function with variables of the type and name of the arguments.
+    """
+    hints = get_type_hints(gen)
+    args = (_var(p.name, hints[p.name]) for p in signature(gen).parameters.values())
+    return gen(*args)
+
+
 ActionLike = Union[Action, BaseExpr]
 
 

python/egglog/egraph.py

@@ -181,7 +181,7 @@ def assert_simplifies(left: BaseExpr, right: BaseExpr) -> None:
         egraph.run(30)
         res = egraph.extract(left)
         print(f"{left} ➡  {res}")
-        egraph.check(eq(right).to(right))
+        egraph.check(eq(left).to(right))
 
 
 assert_simplifies((Term.val(Val(1))).eval(), Val(1))

python/egglog/examples/lambda.py

@@ -0,0 +1,172 @@
+"""
+N-Dimensional Arrays
+====================
+"""
+# mypy: disable-error-code=empty-body
+from __future__ import annotations
+
+from egglog import *
+
+# Story
+# Start with eager arrays
+# Then move onto lazy arrays.
+
+# Then show how the point is we can do this with different packages...
+
+# 1. Wonderful ndarray library with execution on my cpu
+# 2. Come along with that allows generation of code for another platform, thats batched
+# 3. Add cross product, works on both platforms
+# 4. Add specialisation for lazy platform.
+
+# 1. Different people to do different parts, without explicit coordination... No need to ask for permission.
+# 2. Can
+
+
+egraph = EGraph()
+
+
+@egraph.class_
+class Value(BaseExpr):
+    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)),
+)
+
+
+@egraph.class_
+class Values(BaseExpr):
+    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])
+
+
+@egraph.class_
+class NDArray(BaseExpr):
+    """
+    An n-dimensional array.
+    """
+
+    def __getitem__(self, idx: Values) -> Value:
+        ...
+
+    def shape(self) -> Values:
+        ...
+
+
+@egraph.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: BaseExpr, right: BaseExpr) -> 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))
+
+
+@egraph.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)))
+
+
+@egraph.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)))
+
+
+@egraph.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]"))
+
+
+@egraph.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]"))