egglog 7.2.0__cp310-none-win_amd64.whl → 8.0.1__cp310-none-win_amd64.whl

egglog/exp/array_api_loopnest.py

@@ -0,0 +1,149 @@
+"""
+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):
+    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:
+    # 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)
+
+# egraph.register(val.shape)
+# egraph.run(array_api_ruleset.saturate())
+# egraph.extract_multiple(val.shape, 10)
+
+# %%
+
+X = NDArray.var("X")
+assume_shape(X, (3, 2, 3, 4))
+val = linalg_norm(X, (0, 1))
+egraph = EGraph()
+x = egraph.let("x", val.shape[2])
+# egraph.display(n_inline_leaves=0)
+# egraph.extract(x)
+# egraph.saturate(array_api_ruleset, expr=val.shape[2], split_functions=[Int, TRUE, FALSE], n_inline_leaves=2)
+# egraph.run(array_api_ruleset.saturate())
+# egraph.display()
+
+
+# %%

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

@@ -1,8 +1,6 @@
 # mypy: disable-error-code="empty-body"
 from __future__ import annotations
 
-import numpy as np
-
 from egglog import *
 
 from .array_api import *
@@ -13,9 +11,12 @@ from .program_gen import *
 # Depends on `np` as a global variable.
 ##
 
-array_api_program_gen_ruleset = ruleset()
+array_api_program_gen_ruleset = ruleset(name="array_api_program_gen_ruleset")
+array_api_program_gen_eval_ruleset = ruleset(name="array_api_program_gen_eval_ruleset")
 
-array_api_program_gen_schedule = array_api_program_gen_ruleset.saturate() + program_gen_ruleset.saturate()
+array_api_program_gen_schedule = (
+    array_api_program_gen_ruleset | program_gen_ruleset | array_api_program_gen_eval_ruleset | eval_program_rulseset
+).saturate()
 
 
 @function
@@ -56,39 +57,56 @@ 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
 def ndarray_program(x: NDArray) -> Program: ...
 
 
-@function
-def ndarray_function_two(res: NDArray, l: NDArray, r: NDArray) -> Program: ...
+@function(ruleset=array_api_program_gen_ruleset)
+def ndarray_function_two_program(res: NDArray, l: NDArray, r: NDArray) -> Program:
+    return ndarray_program(res).function_two(ndarray_program(l), ndarray_program(r))
 
 
-@array_api_program_gen_ruleset.register
-def _ndarray_function_two(f: Program, res: NDArray, l: NDArray, r: NDArray, o: PyObject):
-    # When we have function, set the program and trigger it to be compiled
-    yield rule(eq(f).to(ndarray_function_two(res, l, r))).then(
-        union(f).with_(ndarray_program(res).function_two(ndarray_program(l), ndarray_program(r))),
-        f.eval_py_object({"np": np}),
-    )
+@function(ruleset=array_api_program_gen_eval_ruleset)
+def ndarray_function_two(res: NDArray, l: NDArray, r: NDArray) -> EvalProgram:
+    return EvalProgram(ndarray_function_two_program(res, l, r), {"np": np})
 
 
 @function
@@ -248,12 +266,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 +303,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/exp/program_gen.py

@@ -83,8 +83,18 @@ class Program(Expr):
         Only keeps the original parent, not any additional ones, so that each set of statements is only added once.
         """
 
-    @method(default=Unit())
-    def eval_py_object(self, globals: object) -> Unit:
+    @property
+    def is_identifer(self) -> Bool:
+        """
+        Returns whether the expression is an identifier. Used so that we don't re-assign any identifiers.
+        """
+
+
+converter(String, Program, Program)
+
+
+class EvalProgram(Expr):
+    def __init__(self, program: Program, globals: object) -> None:
         """
         Evaluates the program and saves as the py_object
         """
@@ -98,38 +108,34 @@ class Program(Expr):
         """
 
     @property
-    def is_identifer(self) -> Bool:
+    def statements(self) -> String:
         """
-        Returns whether the expression is an identifier. Used so that we don't re-assign any identifiers.
+        Returns the statements of the program, if it's been compiled
         """
 
 
-converter(String, Program, Program)
-
-program_gen_ruleset = ruleset()
-
-
-@program_gen_ruleset.register
-def _py_object(p: Program, expr: String, statements: String, g: PyObject):
+@ruleset
+def eval_program_rulseset(ep: EvalProgram, p: Program, expr: String, statements: String, g: PyObject):
     # When we evaluate a program, we first want to compile to a string
-    yield rule(p.eval_py_object(g)).then(p.compile())
+    yield rule(EvalProgram(p, g)).then(p.compile())
     # Then we want to evaluate the statements/expr
     yield rule(
-        p.eval_py_object(g),
+        eq(ep).to(EvalProgram(p, g)),
         eq(p.statements).to(statements),
         eq(p.expr).to(expr),
     ).then(
-        set_(p.py_object).to(
+        set_(ep.py_object).to(
             py_eval(
                 "l['___res']",
                 PyObject.dict(PyObject.from_string("l"), py_exec(join(statements, "\n", "___res = ", expr), g)),
             )
-        )
+        ),
+        set_(ep.statements).to(statements),
     )
 
 
-@program_gen_ruleset.register
-def _compile(
+@ruleset
+def program_gen_ruleset(
     s: String,
     s1: String,
     s2: String,

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/ipython_magic.py

@@ -14,7 +14,7 @@ if IN_IPYTHON:
 
     @needs_local_scope
     @register_cell_magic
-    def egglog(line, cell, local_ns):  # noqa: ANN001, ANN201
+    def egglog(line, cell, local_ns):
         """
         Run an egglog program.