egglog 7.2.0__cp311-none-win_amd64.whl → 8.0.0__cp311-none-win_amd64.whl

egglog/exp/array_api_loopnest.py

@@ -0,0 +1,145 @@
+"""
+In progress module
+
+https://gist.github.com/sklam/5e5737137d48d6e5b816d14a90076f1d
+
+"""
+
+# %%
+# mypy: disable-error-code="empty-body"
+from __future__ import annotations
+
+from egglog import *
+from egglog.exp.array_api import *
+
+
+class ShapeAPI(Expr):
+    @method(unextractable=True)
+    def __init__(self, dims: TupleIntLike) -> None: ...
+
+    @method(unextractable=True)
+    def deselect(self, axis: TupleIntLike) -> ShapeAPI: ...
+
+    @method(unextractable=True)
+    def select(self, axis: TupleIntLike) -> ShapeAPI: ...
+
+    @method(unextractable=True)
+    def to_tuple(self) -> TupleInt: ...
+
+
+@array_api_ruleset.register
+def shape_api_ruleset(dims: TupleInt, axis: TupleInt):  # noqa: ANN201
+    s = ShapeAPI(dims)
+    yield rewrite(s.deselect(axis)).to(
+        ShapeAPI(TupleInt.range(dims.length()).filter(lambda i: ~axis.contains(i)).map(lambda i: dims[i]))
+    )
+    yield rewrite(s.select(axis)).to(
+        ShapeAPI(TupleInt.range(dims.length()).filter(lambda i: axis.contains(i)).map(lambda i: dims[i]))
+    )
+    yield rewrite(s.to_tuple()).to(dims)
+
+
+class OptionalLoopNestAPI(Expr):
+    def __init__(self, value: LoopNestAPI) -> None: ...
+
+    NONE: ClassVar[OptionalLoopNestAPI]
+
+    def unwrap(self) -> LoopNestAPI: ...
+
+
+class LoopNestAPI(Expr):
+    def __init__(self, dim: Int, inner: OptionalLoopNestAPI) -> None: ...
+
+    @classmethod
+    def from_tuple(cls, args: TupleInt) -> OptionalLoopNestAPI: ...
+
+    @method(preserve=True)
+    def __iter__(self) -> Iterator[TupleInt]:
+        return iter(self.indices)
+
+    @property
+    def indices(self) -> TupleTupleInt: ...
+
+    def get_dims(self) -> TupleInt: ...
+
+    def fold(self, fn: Callable[[NDArray, TupleInt], NDArray], init: NDArrayLike) -> NDArray: ...
+
+
+@function
+def tuple_tuple_int_reduce_ndarray(
+    xs: TupleTupleInt, fn: Callable[[NDArray, TupleInt], NDArray], init: NDArray
+) -> NDArray: ...
+
+
+@function
+def tuple_int_map_tuple_int(xs: TupleInt, fn: Callable[[Int], TupleInt]) -> TupleTupleInt: ...
+
+
+@function
+def tuple_tuple_int_product(xs: TupleTupleInt) -> TupleTupleInt: ...
+
+
+@array_api_ruleset.register
+def _loopnest_api_ruleset(
+    head: Int,
+    tail: TupleInt,
+    lna: LoopNestAPI,
+    fn: Callable[[NDArray, TupleInt], NDArray],
+    init: NDArray,
+    dim: Int,
+    idx_fn: Callable[[Int], Int],
+    i: i64,
+):
+    # from_tuple
+    yield rewrite(LoopNestAPI.from_tuple(TupleInt(0, idx_fn))).to(OptionalLoopNestAPI.NONE)
+    yield rewrite(LoopNestAPI.from_tuple(TupleInt(Int(i), idx_fn))).to(
+        OptionalLoopNestAPI(
+            LoopNestAPI(idx_fn(Int(0)), LoopNestAPI.from_tuple(TupleInt(Int(i - 1), lambda i: idx_fn(i + 1))))
+        ),
+        ne(i).to(i64(0)),
+    )
+    # reduce
+    yield rewrite(lna.fold(fn, init)).to(tuple_tuple_int_reduce_ndarray(lna.indices, fn, init))
+    # get_dims
+    yield rewrite(LoopNestAPI(dim, OptionalLoopNestAPI.NONE).get_dims()).to(TupleInt.single(dim))
+    yield rewrite(LoopNestAPI(dim, OptionalLoopNestAPI(lna)).get_dims()).to(TupleInt.single(dim) + lna.get_dims())
+    # indices
+    yield rewrite(lna.indices).to(tuple_tuple_int_product(tuple_int_map_tuple_int(lna.get_dims(), TupleInt.range)))
+
+
+@function(ruleset=array_api_ruleset, unextractable=True)
+def linalg_norm(X: NDArray, axis: TupleIntLike) -> NDArray:  # noqa: N803
+    # peel off the outer shape for result array
+    outshape = ShapeAPI(X.shape).deselect(axis).to_tuple()
+    # get only the inner shape for reduction
+    reduce_axis = ShapeAPI(X.shape).select(axis).to_tuple()
+
+    return NDArray(
+        outshape,
+        X.dtype,
+        lambda k: sqrt(
+            LoopNestAPI.from_tuple(reduce_axis)
+            .unwrap()
+            .fold(lambda carry, i: carry + real(conj(x := X[i + k]) * x), init=0.0)
+        ).to_value(),
+    )
+
+
+# %%
+egraph = EGraph(save_egglog_string=True)
+
+X = NDArray.var("X")
+assume_shape(X, (3, 2, 3, 4))
+val = linalg_norm(X, (0, 1))
+egraph.register(val.shape)
+egraph.run(array_api_ruleset.saturate())
+egraph.extract_multiple(val.shape, 10)
+
+# %%
+egraph = EGraph()
+egraph.register(val.shape[2])
+egraph.run(array_api_ruleset.saturate())
+egraph.display(split_functions=[Int, TRUE, FALSE], n_inline_leaves=2)
+
+
+# %%

egglog/exp/array_api_numba.py

@@ -67,6 +67,6 @@ def _unique_inverse(x: NDArray, i: Int):
     return [
         # Creating a mask array of when the unique inverse is a value is the same as a mask array for when the value is that index of the unique values
         rewrite(unique_inverse(x)[Int(1)] == NDArray.scalar(Value.int(i)), subsume=True).to(
-            x == NDArray.scalar(unique_values(x).index(TupleInt(i)))
+            x == NDArray.scalar(unique_values(x).index((i,)))
         ),
     ]

egglog/exp/array_api_program_gen.py

@@ -56,22 +56,42 @@ def _int_program(i64_: i64, i: Int, j: Int):
 
 
 @function
-def tuple_int_program(x: TupleInt) -> Program: ...
+def tuple_int_program(x: TupleInt) -> Program:
+    ...
+    # Could be rewritten as a fold, but we don't support generic folds yet
+    # return x.fold(Program("("), lambda acc, i: acc + ", " + int_program(i)) + ")"
 
 
 @function
-def tuple_int_program_inner(x: TupleInt) -> Program: ...
+def tuple_int_program_inner(x: TupleInt) -> Program:
+    """
+    Returns the tuple w/ out the parenthesis
+    """
 
 
 @array_api_program_gen_ruleset.register
-def _tuple_int_program(i: Int, j: Int, ti: TupleInt, ti1: TupleInt, ti2: TupleInt):
+def _tuple_int_program(i: Int, ti: TupleInt, k: i64, idx_fn: Callable[[Int], Int], vec_int: Vec[Int]):
     yield rewrite(int_program(ti[i])).to(tuple_int_program(ti) + "[" + int_program(i) + "]")
+    yield rewrite(int_program(ti.length())).to(Program("len(") + tuple_int_program(ti) + ")")
 
     yield rewrite(tuple_int_program(ti)).to(Program("(") + tuple_int_program_inner(ti) + ")")
-    yield rewrite(tuple_int_program_inner(ti1 + ti2)).to(
-        tuple_int_program_inner(ti1) + " " + tuple_int_program_inner(ti2)
+
+    yield rewrite(tuple_int_program_inner(TupleInt(0, idx_fn))).to(Program(""))
+
+    yield rewrite(tuple_int_program_inner(TupleInt(Int(k), idx_fn))).to(
+        int_program(idx_fn(Int(0))) + ", " + tuple_int_program_inner(TupleInt(Int(k - 1), lambda i: idx_fn(i + 1))),
+        ne(k).to(i64(0)),
+    )
+
+    yield rewrite(tuple_int_program_inner(TupleInt.from_vec(Vec[Int]()))).to(Program(""))
+    yield rewrite(tuple_int_program_inner(TupleInt.from_vec(vec_int))).to(
+        int_program(vec_int[0]) + ", " + tuple_int_program_inner(TupleInt.from_vec(vec_int.remove(0))),
+        vec_int.length() > 1,
+    )
+    yield rewrite(tuple_int_program_inner(TupleInt.from_vec(vec_int))).to(
+        int_program(vec_int[0]) + ",",
+        eq(vec_int.length()).to(i64(1)),
     )
-    yield rewrite(tuple_int_program_inner(TupleInt(i))).to(int_program(i) + ",")
 
 
 @function
@@ -248,12 +268,31 @@ def multi_axis_index_key_program(x: MultiAxisIndexKey) -> Program: ...
 
 
 @array_api_program_gen_ruleset.register
-def _multi_axis_index_key_program(l: MultiAxisIndexKey, r: MultiAxisIndexKey, item: MultiAxisIndexKeyItem):
-    yield rewrite(multi_axis_index_key_program(MultiAxisIndexKey(item))).to(multi_axis_index_key_item_program(item))
-    yield rewrite(multi_axis_index_key_program(l + r)).to(
-        multi_axis_index_key_program(l) + ", " + multi_axis_index_key_program(r)
+def _multi_axis_index_key_program(
+    idx_fn: Callable[[Int], MultiAxisIndexKeyItem], k: i64, vec: Vec[MultiAxisIndexKeyItem], i: MultiAxisIndexKeyItem
+):
+    yield rewrite(multi_axis_index_key_program(MultiAxisIndexKey(0, idx_fn))).to(Program(""))
+
+    yield rewrite(multi_axis_index_key_program(MultiAxisIndexKey(Int(k), idx_fn))).to(
+        multi_axis_index_key_item_program(idx_fn(Int(0)))
+        + ", "
+        + multi_axis_index_key_program(MultiAxisIndexKey(Int(k - 1), lambda i: idx_fn(i + 1))),
+        ne(k).to(i64(0)),
+    )
+
+    yield rewrite(multi_axis_index_key_program(MultiAxisIndexKey.from_vec(Vec[MultiAxisIndexKeyItem]()))).to(
+        Program("")
+    )
+    yield rewrite(multi_axis_index_key_program(MultiAxisIndexKey.from_vec(vec))).to(
+        multi_axis_index_key_item_program(vec[0]) + ",",
+        eq(vec.length()).to(i64(1)),
+    )
+    yield rewrite(multi_axis_index_key_program(MultiAxisIndexKey.from_vec(vec))).to(
+        multi_axis_index_key_item_program(vec[0])
+        + ", "
+        + multi_axis_index_key_program(MultiAxisIndexKey.from_vec(vec.remove(0))),
+        vec.length() > 1,
     )
-    yield rewrite(multi_axis_index_key_program(MultiAxisIndexKey.EMPTY)).to(Program("()"))
 
 
 @function
@@ -266,7 +305,7 @@ def _index_key_program(i: Int, s: Slice, key: MultiAxisIndexKey, a: NDArray):
     yield rewrite(index_key_program(IndexKey.int(i))).to(int_program(i))
     yield rewrite(index_key_program(IndexKey.slice(s))).to(slice_program(s))
     yield rewrite(index_key_program(IndexKey.multi_axis(key))).to(multi_axis_index_key_program(key))
-    yield rewrite(index_key_program(ndarray_index(a))).to(ndarray_program(a))
+    yield rewrite(index_key_program(IndexKey.ndarray(a))).to(ndarray_program(a))
 
 
 @function

egglog/functionalize.py

@@ -0,0 +1,91 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from functools import partial
+from inspect import Parameter, signature
+from typing import Any, TypeVar, cast
+
+__all__ = ["functionalize"]
+
+
+T = TypeVar("T", bound=Callable)
+
+
+# TODO: Add `to_lift` param so that we only transform those with vars in them to args
+
+
+def functionalize(f: T, get_annotation: Callable[[object], type | None]) -> T:
+    """
+    Takes a function and returns a new function with all names (co_names) and free variables (co_freevars) added as arguments
+    and then partially applied with their values. The second arg, get_annotation, will be applied to all values
+    to get their type annotation. If it is None, that arg will not be added as a parameter.
+
+    For example if you have:
+
+        def get_annotation(x): return int if x <= 10 else None
+
+        g = 10
+        def f(a, a2):
+            def h(b: Z):
+                return a + a2 + b + g
+
+            return functionalize(h, get_annotation)
+        res = f(9, 11)
+
+    It should be equivalent to (according to body, signature, and annotations) (Note that the new arguments will be positional only):
+
+        def h(a: get_annotation(a), g: get_annotation(g), b: Z):
+            return a + b + g
+        res = partial(h, a, g)
+    """
+    code = f.__code__
+    names = tuple(n for n in code.co_names if n in f.__globals__)
+    free_vars = code.co_freevars
+
+    global_values: list[Any] = [f.__globals__[name] for name in names]
+    free_var_values = [cell.cell_contents for cell in f.__closure__] if f.__closure__ else []
+    assert len(free_var_values) == len(free_vars), "Free vars and their values do not match"
+    global_values_filtered = [
+        (i, name, value, annotation)
+        for i, (name, value) in enumerate(zip(names, global_values, strict=True))
+        if (annotation := get_annotation(value)) is not None
+    ]
+    free_var_values_filtered = [
+        (i, name, value, annotation)
+        for i, (name, value) in enumerate(zip(free_vars, free_var_values, strict=True))
+        if (annotation := get_annotation(value)) is not None
+    ]
+    additional_arg_filtered = global_values_filtered + free_var_values_filtered
+
+    # Create a wrapper function
+    def wrapper(*args):
+        # Split args into names, free vars and other args
+        name_args, free_var_args, rest_args = (
+            args[: (n_names := len(global_values_filtered))],
+            args[n_names : (n_args := len(additional_arg_filtered))],
+            args[n_args:],
+        )
+        # Update globals with names
+        f.__globals__.update({
+            name: arg for (_, name, _, _), arg in zip(global_values_filtered, name_args, strict=False)
+        })
+        # update function free vars with free var args
+        for (i, _, _, _), value in zip(free_var_values_filtered, free_var_args, strict=True):
+            assert f.__closure__, "Function does not have closure"
+            f.__closure__[i].cell_contents = value
+        return f(*rest_args)
+
+    # Set the signature of the new function to a signature with the free vars and names added as arguments
+    orig_signature = signature(f)
+    wrapper.__signature__ = orig_signature.replace(  # type: ignore[attr-defined]
+        parameters=[
+            *[Parameter(n, Parameter.POSITIONAL_OR_KEYWORD) for _, n, _, _ in additional_arg_filtered],
+            *orig_signature.parameters.values(),
+        ]
+    )
+    # Set the annotations of the new function to the annotations of the original function + annotations of passed in values
+    wrapper.__annotations__ = f.__annotations__ | {n: a for _, n, _, a in additional_arg_filtered}
+    wrapper.__name__ = f.__name__
+
+    # Partially apply the wrapper function with the current values of the free vars
+    return cast(T, partial(wrapper, *(v for _, _, v, _ in additional_arg_filtered)))

egglog/pretty.py

@@ -16,6 +16,7 @@ from .declarations import *
 if TYPE_CHECKING:
     from collections.abc import Mapping
 
+
 __all__ = [
     "pretty_decl",
     "pretty_callable_ref",
@@ -77,9 +78,9 @@ def pretty_decl(
 
     This will use re-format the result and put the expression on the last line, preceeded by the statements.
     """
-    traverse = TraverseContext()
+    traverse = TraverseContext(decls)
     traverse(decl, toplevel=True)
-    pretty = traverse.pretty(decls)
+    pretty = traverse.pretty()
     expr = pretty(decl, ruleset_name=ruleset_name)
     if wrapping_fn:
         expr = f"{wrapping_fn}({expr})"
@@ -106,15 +107,20 @@ def pretty_callable_ref(
     """
     # Pass in three dummy args, which are the max used for any operation that
     # is not a generic function call
-    args: list[ExprDecl] = [VarDecl(ARG_STR)] * 3
+    args: list[ExprDecl] = [VarDecl(ARG_STR, False)] * 3
     if first_arg:
         args.insert(0, first_arg)
-    res = PrettyContext(decls, defaultdict(lambda: 0))._call_inner(
-        ref, args, bound_tp_params=bound_tp_params, parens=False
-    )
+    context = PrettyContext(decls, defaultdict(lambda: 0))
+    res = context._call_inner(ref, args, bound_tp_params=bound_tp_params, parens=False)
     # Either returns a function or a function with args. If args are provided, they would just be called,
     # on the function, so return them, because they are dummies
-    return res[0] if isinstance(res, tuple) else res
+    if isinstance(res, tuple):
+        name = res[0]
+        # if this is an unnamed function, return it but don't partially apply any args
+        if isinstance(name, UnnamedFunctionRef):
+            return context._pretty_function_body(name, [])
+        return name
+    return res
 
 
 # TODO: Add a different pretty callable ref that doesnt fill in wholes but instead returns the function
@@ -128,16 +134,18 @@ class TraverseContext:
     expression has.
     """
 
+    decls: Declarations
+
     # All expressions we have seen (incremented the parent counts of all children)
     _seen: set[AllDecls] = field(default_factory=set)
     # The number of parents for each expressions
     parents: Counter[AllDecls] = field(default_factory=Counter)
 
-    def pretty(self, decls: Declarations) -> PrettyContext:
+    def pretty(self) -> PrettyContext:
         """
         Create a pretty context from the state of this traverse context.
         """
-        return PrettyContext(decls, self.parents)
+        return PrettyContext(self.decls, self.parents)
 
     def __call__(self, decl: AllDecls, toplevel: bool = False) -> None:  # noqa: C901
         if not toplevel:
@@ -169,9 +177,13 @@ class TraverseContext:
                     if isinstance(de, DefaultRewriteDecl):
                         continue
                     self(de)
-            case CallDecl(_, exprs, _):
-                for e in exprs:
-                    self(e.expr)
+            case CallDecl(ref, exprs, _):
+                match ref:
+                    case FunctionRef(UnnamedFunctionRef(_, res)):
+                        self(res.expr)
+                    case _:
+                        for e in exprs:
+                            self(e.expr)
             case RunDecl(_, until):
                 if until:
                     for f in until:
@@ -244,8 +256,7 @@ class PrettyContext:
             case CallDecl(_, _, _):
                 return self._call(decl, parens)
             case PartialCallDecl(CallDecl(ref, typed_args, _)):
-                arg_strs = (_pretty_callable(ref), *(self(a.expr, parens=False, unwrap_lit=True) for a in typed_args))
-                return f"UnstableFn({', '.join(arg_strs)})", "fn"
+                return self._pretty_partial(ref, [a.expr for a in typed_args]), "fn"
             case PyObjectDecl(value):
                 return repr(value) if unwrap_lit else f"PyObject({value!r})", "PyObject"
             case ActionCommandDecl(action):
@@ -352,12 +363,16 @@ class PrettyContext:
             has_multiple_parents = self.parents[first_arg] > 1
             self.names[decl] = expr_name = self._name_expr(tp_name, expr_str, copy_identifier=has_multiple_parents)
             # Set the first arg to be the name of the mutated arg and return the name
-            args[0] = VarDecl(expr_name)
+            args[0] = VarDecl(expr_name, True)
         else:
             expr_name = None
         res = self._call_inner(ref, args, decl.bound_tp_params, parens)
         expr = (
-            f"{res[0]}({', '.join(self(a, parens=False, unwrap_lit=True) for a in res[1])})"
+            (
+                f"{name}({', '.join(self(a, parens=False, unwrap_lit=True) for a in res[1])})"
+                if isinstance((name := res[0]), str)
+                else ((called := self._pretty_function_body(name, res[1])) if not parens else f"({called})")
+            )
             if isinstance(res, tuple)
             else res
         )
@@ -368,8 +383,12 @@ class PrettyContext:
         return expr, tp_name
 
     def _call_inner(  # noqa: PLR0911
-        self, ref: CallableRef, args: list[ExprDecl], bound_tp_params: tuple[JustTypeRef, ...] | None, parens: bool
-    ) -> tuple[str, list[ExprDecl]] | str:
+        self,
+        ref: CallableRef,
+        args: list[ExprDecl],
+        bound_tp_params: tuple[JustTypeRef, ...] | None,
+        parens: bool,
+    ) -> tuple[str | UnnamedFunctionRef, list[ExprDecl]] | str:
         """
         Pretty print the call, returning either the full function call or a tuple of the function and the args.
         """
@@ -408,6 +427,8 @@ class PrettyContext:
             case InitRef(class_name):
                 tp_ref = JustTypeRef(class_name, bound_tp_params or ())
                 return str(tp_ref), args
+            case UnnamedFunctionRef():
+                return ref, args
         assert_never(ref)
 
     def _generate_name(self, typ: str) -> str:
@@ -428,29 +449,52 @@ class PrettyContext:
             self.statements.append(f"{name} = {expr_str}")
         return name
 
+    def _pretty_partial(self, ref: CallableRef, args: list[ExprDecl]) -> str:
+        """
+        Returns a partial function call as a string.
+        """
+        match ref:
+            case FunctionRef(name):
+                fn = name
+            case UnnamedFunctionRef():
+                return self._pretty_function_body(ref, args)
+            case (
+                ClassMethodRef(class_name, method_name)
+                | MethodRef(class_name, method_name)
+                | PropertyRef(class_name, method_name)
+            ):
+                fn = f"{class_name}.{method_name}"
+            case InitRef(class_name):
+                fn = class_name
+            case ConstantRef(_):
+                msg = "Constants should not be callable"
+                raise NotImplementedError(msg)
+            case ClassVariableRef(_, _):
+                msg = "Class variables should not be callable"
+                raise NotADirectoryError(msg)
+            case _:
+                assert_never(ref)
+        if not args:
+            return fn
+        arg_strs = (
+            fn,
+            *(self(a, parens=False, unwrap_lit=True) for a in args),
+        )
+        return f"partial({', '.join(arg_strs)})"
 
-def _pretty_callable(ref: CallableRef) -> str:
-    """
-    Returns a function call as a string.
-    """
-    match ref:
-        case FunctionRef(name):
-            return name
-        case (
-            ClassMethodRef(class_name, method_name)
-            | MethodRef(class_name, method_name)
-            | PropertyRef(class_name, method_name)
-        ):
-            return f"{class_name}.{method_name}"
-        case InitRef(class_name):
-            return class_name
-        case ConstantRef(_):
-            msg = "Constants should not be callable"
-            raise NotImplementedError(msg)
-        case ClassVariableRef(_, _):
-            msg = "Class variables should not be callable"
-            raise NotADirectoryError(msg)
-    assert_never(ref)
+    def _pretty_function_body(self, fn: UnnamedFunctionRef, args: list[ExprDecl]) -> str:
+        """
+        Pretty print the body of a function, partially applying some arguments.
+        """
+        var_args = fn.args
+        replacements = {var_arg: TypedExprDecl(var_arg.tp, arg) for var_arg, arg in zip(var_args, args, strict=False)}
+        var_args = var_args[len(args) :]
+        res = replace_typed_expr(fn.res, replacements)
+        arg_names = fn.args[len(args) :]
+        prefix = "lambda"
+        if arg_names:
+            prefix += f" {', '.join(self(a.expr) for a in arg_names)}"
+        return f"{prefix}: {self(res.expr)}"
 
 
 def _plot_line_length(expr: object):  # pragma: no cover