High level Python API

[saulshanabrook]

Hi!

I wanted to open an issue to let folks know I have a version of a high level API wrapper working in Python. I published a new version of the Python package and put together a tutorial showing how to re-create the matrix examples in Python:

from egg_smol import *

egraph = EGraph()

@egraph.class_
class Dim(BaseExpr):
    """
    A dimension of a matix.

    >>> Dim(3) * Dim.named("n")
    Dim(3) * Dim.named("n")
    """
    def __init__(self, value: i64Like) -> None:
        ...

    @classmethod
    def named(cls, name: StringLike) -> Dim:
        ...

    def __mul__(self, other: Dim) -> Dim:
        ...

a, b, c = vars_("a b c", 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),
)

x = Dim.named("x")
ten = Dim(10)
res = x * ten * ten

graph.register(let("", res))
egraph.run(10)
egraph.extract(res)

It also exposes the rust API more directly, which the high level API is built on top of.

A couple of fun things that you get with the high level API, is static type analysis provided by any of Python's type checkers (mypy, pylance, etc), to verify that the types (aka sorts) are correct before even running the program. Also, instead of just having everything be a function, we can use Python's operator overloading so that things like a * b are translated into Add(a, b).

This requires knowing all the types at compile time, so that we can differentiate the operator overloading based on the argument types, which I have implemented on the Python side.

It also exposes a few oddnesses with at least the version of egg-smol that I am using, like the use of let as a way of adding expressions to the e-graph, which I use in the Python tutorial and is also used in the matrix example in egg-smol.

I am currently pinned to an older version of egg-smol, so I can focus on getting this API worked out. I do plan on upgrading the underlying dependency next and adding support for the new features.

I would also like to try to add on some graph visualization next, to help understand the state of the e-graph visually as things progress...

I would love any feedback on the Python API! I had to make some design decisions on how to expose things and would love to hear how they land with people. Although it's definitely more verbose than the lisp API, there are some advantages of being a DSL in a host language with tooling as I mentioned above.

Thanks for all the work on this library! It's very exciting to be getting closer to a language-independent way of expressing DSLs and translations between them :)

[mwillsey]

Very cool! I think Python bindings are eventually going to be very valuable. I like the idea of sticking pretty close to the SMT / z3 model. If anything I would want to stick as close as possible to that.

The static typing is definitely neat! But I worry that it basically requires a static re-implementation of egg's internal typechecker.

[saulshanabrook]

But I worry that it basically requires a static re-implementation of egg's internal typechecker.

Yes it does. I don't see a way around that if we want to be able to support operator overloading... It isn't such a huge lift though. I mean especially now, the type system is very simple, but yeah eventually it will expand.

Another option would be for the Python API to call out the typechecking API in Rust, as it builds up expressions. That could be helpful if/when the type system gets more complicated?

One difference is that in the Python world all variables are typed, whereas in egg-smol they are untyped. The typing makes the analysis a bit easier as well from the Python side of things.

Had you thought about having variables be typed in egg-smol?

EDIT: Oh I just wanted to mention, to make sure we are on the same page, the only reason we need to do any type analysis at all at the Python level is to know what types we have so we know what methods are called. I.e. the difference between say MyExpr.__add__ and i64.__add__. So those can be mapped to different egg-smol functions.

The nice thing is we don't need to implement the actual static checking, because that is already handled by Mypy etc since we conform to normal Python objects/functions (as far as it knows at least. Of course at runtime they aren't, but statically we pretend they are).

I like the idea of sticking pretty close to the SMT / z3 model. If anything I would want to stick as close as possible to that.

I'll take a look at the z3 API. I haven't used it much, but I can see what we can copy from there!

One thing that would be nice would be a higher level simplify command, which adds an expression to the e-graph, runs some number of steps, and then extracts it. That is currently a bit laborious.

[mwillsey]

Variables in egg-smol are typed, but the types are inferred. We should probably add syntax for explicit type annotation as well.

…

On Tue, Mar 28, 2023 at 12:07 PM Saul Shanabrook ***@***.***> wrote: But I worry that it basically requires a static re-implementation of egg's internal typechecker. Yes it does. I don't see a way around that if we want to be able to support operator overloading... It isn't such a huge lift though. I mean especially now, the type system is very simple, but yeah eventually it will expand. Another option would be for the Python API to call out the typechecking API in Rust, as it builds up expressions. That could be helpful if/when the type system gets more complicated? One difference is that in the Python world all variables are typed, whereas in egg-smol they are untyped. The typing makes the analysis a bit easier as well from the Python side of things. Had you thought about having variables be typed in egg-smol? I like the idea of sticking pretty close to the SMT / z3 model. If anything I would want to stick as close as possible to that. I'll take a look at the z3 API <https://ericpony.github.io/z3py-tutorial/guide-examples.htm>. I haven't used it much, but I can see what we can copy from there! One thing that would be nice would be a higher level simplify command, which adds an expression to the e-graph, runs some number of steps, and then extracts it. That is currently a bit laborious. — Reply to this email directly, view it on GitHub <#126 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AANTPTEOKYIDSHR2DXAIXLDW6MZIJANCNFSM6AAAAAAWJ3JDNQ> . You are receiving this because you commented.Message ID: ***@***.***>