egglog 12.0.0__cp313-cp313t-manylinux_2_17_ppc64.manylinux2014_ppc64.whl

egglog/exp/array_api_jit.py

@@ -0,0 +1,51 @@
+import inspect
+from collections.abc import Callable
+from typing import TypeVar, cast
+
+import numpy as np
+
+from egglog import EGraph, greedy_dag_cost_model
+from egglog.exp.array_api import NDArray, set_array_api_egraph, try_evaling
+from egglog.exp.array_api_numba import array_api_numba_schedule
+from egglog.exp.array_api_program_gen import EvalProgram, array_api_program_gen_schedule, ndarray_function_two_program
+
+from .program_gen import Program
+
+X = TypeVar("X", bound=Callable)
+
+
+def jit(
+    fn: X,
+    *,
+    handle_expr: Callable[[NDArray], None] | None = None,
+    handle_optimized_expr: Callable[[NDArray], None] | None = None,
+) -> X:
+    """
+    Jit compiles a function
+    """
+    egraph, res, res_optimized, program = function_to_program(fn, save_egglog_string=False)
+    if handle_expr:
+        handle_expr(res)
+    if handle_optimized_expr:
+        handle_optimized_expr(res_optimized)
+    fn_program = EvalProgram(program, {"np": np})
+    return cast("X", try_evaling(egraph, array_api_program_gen_schedule, fn_program, fn_program.as_py_object))
+
+
+def function_to_program(fn: Callable, save_egglog_string: bool) -> tuple[EGraph, NDArray, NDArray, Program]:
+    sig = inspect.signature(fn)
+    arg1, arg2 = sig.parameters.keys()
+    egraph = EGraph(save_egglog_string=save_egglog_string)
+    with egraph:
+        with set_array_api_egraph(egraph):
+            res = fn(NDArray.var(arg1), NDArray.var(arg2))
+        egraph.register(res)
+        egraph.run(array_api_numba_schedule)
+        res_optimized = egraph.extract(res, cost_model=greedy_dag_cost_model())
+
+    return (
+        egraph,
+        res,
+        res_optimized,
+        ndarray_function_two_program(res_optimized, NDArray.var(arg1), NDArray.var(arg2)),
+    )

egglog/exp/array_api_loopnest.py

@@ -0,0 +1,74 @@
+"""
+Example module to replicate behavior expressed in
+
+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 *
+
+__all__ = ["LoopNestAPI", "OptionalLoopNestAPI", "ShapeAPI"]
+
+
+class ShapeAPI(Expr):
+    def __init__(self, dims: TupleIntLike) -> None: ...
+
+    def deselect(self, axis: TupleIntLike) -> ShapeAPI: ...
+
+    def select(self, axis: TupleIntLike) -> ShapeAPI: ...
+
+    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), subsume=True).to(
+        ShapeAPI(TupleInt.range(dims.length()).filter(lambda i: ~axis.contains(i)).map(lambda i: dims[i]))
+    )
+    yield rewrite(s.select(axis), subsume=True).to(
+        ShapeAPI(TupleInt.range(dims.length()).filter(lambda i: axis.contains(i)).map(lambda i: dims[i]))
+    )
+    yield rewrite(s.to_tuple(), subsume=True).to(dims)
+
+
+class OptionalLoopNestAPI(Expr):
+    def __init__(self, value: LoopNestAPI) -> None: ...
+
+    NONE: ClassVar[OptionalLoopNestAPI]
+
+    def unwrap(self) -> LoopNestAPI: ...
+
+
+class LoopNestAPI(Expr, ruleset=array_api_ruleset):
+    def __init__(self, last: Int, before: OptionalLoopNestAPI) -> None: ...
+
+    @classmethod
+    def from_tuple(cls, args: TupleIntLike) -> OptionalLoopNestAPI: ...
+
+    @method(preserve=True)
+    def __iter__(self) -> Iterator[TupleInt]:
+        return iter(self.indices())
+
+    def indices(self) -> TupleTupleInt:
+        return self.get_dims().map_tuple_int(TupleInt.range).product()
+
+    def get_dims(self) -> TupleInt: ...
+
+
+@array_api_ruleset.register
+def _loopnest_api_ruleset(lna: LoopNestAPI, dim: Int, ti: TupleInt, idx_fn: Callable[[Int], Int], i: i64):
+    # from_tuple
+    yield rewrite(LoopNestAPI.from_tuple(TupleInt.EMPTY), subsume=True).to(OptionalLoopNestAPI.NONE)
+    yield rewrite(LoopNestAPI.from_tuple(ti.append(dim)), subsume=True).to(
+        OptionalLoopNestAPI(LoopNestAPI(dim, LoopNestAPI.from_tuple(ti)))
+    )
+    # get_dims
+    yield rewrite(LoopNestAPI(dim, OptionalLoopNestAPI.NONE).get_dims(), subsume=True).to(TupleInt.single(dim))
+    yield rewrite(LoopNestAPI(dim, OptionalLoopNestAPI(lna)).get_dims(), subsume=True).to(lna.get_dims().append(dim))
+    # unwrap
+    yield rewrite(OptionalLoopNestAPI(lna).unwrap()).to(lna)

egglog/exp/array_api_numba.py

@@ -0,0 +1,69 @@
+# mypy: disable-error-code="empty-body"
+"""
+Module for generating array api code that works with Numba.
+"""
+
+from __future__ import annotations
+
+from egglog import *
+from egglog.exp.array_api import *
+
+array_api_numba_ruleset = ruleset()
+array_api_numba_schedule = (array_api_combined_ruleset | array_api_numba_ruleset).saturate()
+# For these rules, we not only wanna rewrite, we also want to subsume the original expression,
+# so that the rewritten one is used, even if the original one is simpler.
+
+
+# Rewrite mean(x, <int>, <expand dims>) to use sum b/c numba cant do mean with axis
+# https://github.com/numba/numba/issues/1269
+@array_api_numba_ruleset.register
+def _mean(y: NDArray, x: NDArray, i: Int):
+    axis = OptionalIntOrTuple.some(IntOrTuple.int(i))
+    res = sum(x, axis) / NDArray.scalar(Value.int(x.shape[i]))
+
+    yield rewrite(mean(x, axis, FALSE), subsume=True).to(res)
+    yield rewrite(mean(x, axis, TRUE), subsume=True).to(expand_dims(res, i))
+
+
+# Rewrite std(x, <int>) to use mean and sum b/c numba cant do std with axis
+@array_api_numba_ruleset.register
+def _std(y: NDArray, x: NDArray, i: Int):
+    axis = OptionalIntOrTuple.some(IntOrTuple.int(i))
+    # https://numpy.org/doc/stable/reference/generated/numpy.std.html
+    # "std = sqrt(mean(x)), where x = abs(a - a.mean())**2."
+    yield rewrite(
+        std(x, axis),
+        subsume=True,
+    ).to(
+        sqrt(mean(square(x - mean(x, axis, keepdims=TRUE)), axis)),
+    )
+
+
+# rewrite unique_counts to count each value one by one, since numba doesn't support np.unique(..., return_counts=True)
+@function(ruleset=array_api_numba_ruleset, subsume=True)
+def count_values(x: NDArrayLike, values: TupleValueLike) -> TupleValue:
+    """
+    Returns a tuple of the count of each of the values in the array as values
+    """
+    x = cast(NDArray, x)
+    values = cast(TupleValue, values)
+    return TupleValue(values.length(), lambda i: sum(x == values[i]).to_value())
+
+
+@array_api_numba_ruleset.register
+def _unique_counts(x: NDArray, c: NDArray, tv: TupleValue, v: Value):
+    return [
+        # The unique counts are the count of all the unique values
+        rewrite(unique_counts(x)[1], subsume=True).to(NDArray.vector(count_values(x, unique_values(x).to_values()))),
+    ]
+
+
+# do the same for unique_inverse
+@array_api_numba_ruleset.register
+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((i,)))
+        ),
+    ]