egglog 6.1.0__cp311-none-win_amd64.whl → 7.1.0__cp311-none-win_amd64.whl

egglog/exp/array_api.py

@@ -18,13 +18,13 @@ from egglog.runtime import RuntimeExpr
 from .program_gen import *
 
 if TYPE_CHECKING:
-    from collections.abc import Iterator
+    from collections.abc import Callable, Iterator
     from types import ModuleType
 
 # Pretend that exprs are numbers b/c sklearn does isinstance checks
 numbers.Integral.register(RuntimeExpr)
 
-array_api_ruleset = ruleset()
+array_api_ruleset = ruleset(name="array_api_ruleset")
 array_api_schedule = array_api_ruleset.saturate()
 
 
@@ -36,10 +36,14 @@ class Boolean(Expr):
     @property
     def bool(self) -> Bool: ...
 
-    def __or__(self, other: Boolean) -> Boolean: ...
+    def __or__(self, other: BooleanLike) -> Boolean: ...
 
-    def __and__(self, other: Boolean) -> Boolean: ...
+    def __and__(self, other: BooleanLike) -> Boolean: ...
 
+    def if_int(self, true_value: Int, false_value: Int) -> Int: ...
+
+
+BooleanLike = Boolean | bool
 
 TRUE = constant("TRUE", Boolean)
 FALSE = constant("FALSE", Boolean)
@@ -47,7 +51,7 @@ converter(bool, Boolean, lambda x: TRUE if x else FALSE)
 
 
 @array_api_ruleset.register
-def _bool(x: Boolean):
+def _bool(x: Boolean, i: Int, j: Int):
     return [
         rule(eq(x).to(TRUE)).then(set_(x.bool).to(Bool(True))),
         rule(eq(x).to(FALSE)).then(set_(x.bool).to(Bool(False))),
@@ -55,82 +59,8 @@ def _bool(x: Boolean):
         rewrite(FALSE | x).to(x),
         rewrite(TRUE & x).to(x),
         rewrite(FALSE & x).to(FALSE),
-    ]
-
-
-class DType(Expr):
-    float64: ClassVar[DType]
-    float32: ClassVar[DType]
-    int64: ClassVar[DType]
-    int32: ClassVar[DType]
-    object: ClassVar[DType]
-    bool: ClassVar[DType]
-
-    def __eq__(self, other: DType) -> Boolean:  # type: ignore[override]
-        ...
-
-
-float64 = DType.float64
-float32 = DType.float32
-int32 = DType.int32
-int64 = DType.int64
-
-_DTYPES = [float64, float32, int32, int64, DType.object]
-
-converter(type, DType, lambda x: convert(np.dtype(x), DType))
-converter(type(np.dtype), DType, lambda x: getattr(DType, x.name))  # type: ignore[call-overload]
-array_api_ruleset.register(
-    *(rewrite(l == r).to(TRUE if l is r else FALSE) for l, r in itertools.product(_DTYPES, repeat=2))
-)
-
-
-class IsDtypeKind(Expr):
-    NULL: ClassVar[IsDtypeKind]
-
-    @classmethod
-    def string(cls, s: StringLike) -> IsDtypeKind: ...
-
-    @classmethod
-    def dtype(cls, d: DType) -> IsDtypeKind: ...
-
-    @method(cost=10)
-    def __or__(self, other: IsDtypeKind) -> IsDtypeKind: ...
-
-
-# TODO: Make kind more generic to support tuples.
-@function
-def isdtype(dtype: DType, kind: IsDtypeKind) -> Boolean: ...
-
-
-converter(DType, IsDtypeKind, lambda x: IsDtypeKind.dtype(x))
-converter(str, IsDtypeKind, lambda x: IsDtypeKind.string(x))
-converter(
-    tuple, IsDtypeKind, lambda x: convert(x[0], IsDtypeKind) | convert(x[1:], IsDtypeKind) if x else IsDtypeKind.NULL
-)
-
-
-@array_api_ruleset.register
-def _isdtype(d: DType, k1: IsDtypeKind, k2: IsDtypeKind):
-    return [
-        rewrite(isdtype(DType.float32, IsDtypeKind.string("integral"))).to(FALSE),
-        rewrite(isdtype(DType.float64, IsDtypeKind.string("integral"))).to(FALSE),
-        rewrite(isdtype(DType.object, IsDtypeKind.string("integral"))).to(FALSE),
-        rewrite(isdtype(DType.int64, IsDtypeKind.string("integral"))).to(TRUE),
-        rewrite(isdtype(DType.int32, IsDtypeKind.string("integral"))).to(TRUE),
-        rewrite(isdtype(DType.float32, IsDtypeKind.string("real floating"))).to(TRUE),
-        rewrite(isdtype(DType.float64, IsDtypeKind.string("real floating"))).to(TRUE),
-        rewrite(isdtype(DType.object, IsDtypeKind.string("real floating"))).to(FALSE),
-        rewrite(isdtype(DType.int64, IsDtypeKind.string("real floating"))).to(FALSE),
-        rewrite(isdtype(DType.int32, IsDtypeKind.string("real floating"))).to(FALSE),
-        rewrite(isdtype(DType.float32, IsDtypeKind.string("complex floating"))).to(FALSE),
-        rewrite(isdtype(DType.float64, IsDtypeKind.string("complex floating"))).to(FALSE),
-        rewrite(isdtype(DType.object, IsDtypeKind.string("complex floating"))).to(FALSE),
-        rewrite(isdtype(DType.int64, IsDtypeKind.string("complex floating"))).to(FALSE),
-        rewrite(isdtype(DType.int32, IsDtypeKind.string("complex floating"))).to(FALSE),
-        rewrite(isdtype(d, IsDtypeKind.NULL)).to(FALSE),
-        rewrite(isdtype(d, IsDtypeKind.dtype(d))).to(TRUE),
-        rewrite(isdtype(d, k1 | k2)).to(isdtype(d, k1) | isdtype(d, k2)),
-        rewrite(k1 | IsDtypeKind.NULL).to(k1),
+        rewrite(TRUE.if_int(i, j)).to(i),
+        rewrite(FALSE.if_int(i, j)).to(j),
     ]
 
 
@@ -264,10 +194,13 @@ converter(int, Int, lambda x: Int(x))
 
 
 class Float(Expr):
+    # Differentiate costs of three constructors so extraction is deterministic if all three are present
+    @method(cost=3)
     def __init__(self, value: f64Like) -> None: ...
 
     def abs(self) -> Float: ...
 
+    @method(cost=2)
     @classmethod
     def rational(cls, r: Rational) -> Float: ...
 
@@ -324,7 +257,7 @@ class TupleInt(Expr):
 
     def __getitem__(self, i: Int) -> Int: ...
 
-    def product(self) -> Int: ...
+    def fold(self, init: Int, f: Callable[[Int, Int], Int]) -> Int: ...
 
 
 converter(
@@ -339,7 +272,7 @@ converter(
 
 
 @array_api_ruleset.register
-def _tuple_int(ti: TupleInt, ti2: TupleInt, i: Int, i2: Int, k: i64):
+def _tuple_int(ti: TupleInt, ti2: TupleInt, i: Int, i2: Int, k: i64, f: Callable[[Int, Int], Int]):
     return [
         rewrite(ti + TupleInt.EMPTY).to(ti),
         rewrite(TupleInt(i).length()).to(Int(1)),
@@ -348,10 +281,10 @@ def _tuple_int(ti: TupleInt, ti2: TupleInt, i: Int, i2: Int, k: i64):
         rewrite((TupleInt(i) + ti)[Int(0)]).to(i),
         # Rule for indexing > 0
         rule(eq(i).to((TupleInt(i2) + ti)[Int(k)]), k > 0).then(union(i).with_(ti[Int(k - 1)])),
-        # Product
-        rewrite(TupleInt(i).product()).to(i),
-        rewrite((TupleInt(i) + ti).product()).to(i * ti.product()),
-        rewrite(TupleInt.EMPTY.product()).to(Int(1)),
+        # fold
+        rewrite(TupleInt.EMPTY.fold(i, f)).to(i),
+        rewrite(TupleInt(i2).fold(i, f)).to(f(i, i2)),
+        rewrite((TupleInt(i2) + ti).fold(i, f)).to(ti.fold(f(i, i2), f)),
     ]
 
 
@@ -366,6 +299,85 @@ converter(type(None), OptionalInt, lambda _: OptionalInt.none)
 converter(Int, OptionalInt, OptionalInt.some)
 
 
+class DType(Expr):
+    float64: ClassVar[DType]
+    float32: ClassVar[DType]
+    int64: ClassVar[DType]
+    int32: ClassVar[DType]
+    object: ClassVar[DType]
+    bool: ClassVar[DType]
+
+    def __eq__(self, other: DType) -> Boolean:  # type: ignore[override]
+        ...
+
+
+float64 = DType.float64
+float32 = DType.float32
+int32 = DType.int32
+int64 = DType.int64
+
+_DTYPES = [float64, float32, int32, int64, DType.object]
+
+converter(type, DType, lambda x: convert(np.dtype(x), DType))
+converter(type(np.dtype), DType, lambda x: getattr(DType, x.name))  # type: ignore[call-overload]
+
+
+@array_api_ruleset.register
+def _():
+    for l, r in itertools.product(_DTYPES, repeat=2):
+        yield rewrite(l == r).to(TRUE if l is r else FALSE)
+
+
+class IsDtypeKind(Expr):
+    NULL: ClassVar[IsDtypeKind]
+
+    @classmethod
+    def string(cls, s: StringLike) -> IsDtypeKind: ...
+
+    @classmethod
+    def dtype(cls, d: DType) -> IsDtypeKind: ...
+
+    @method(cost=10)
+    def __or__(self, other: IsDtypeKind) -> IsDtypeKind: ...
+
+
+# TODO: Make kind more generic to support tuples.
+@function
+def isdtype(dtype: DType, kind: IsDtypeKind) -> Boolean: ...
+
+
+converter(DType, IsDtypeKind, lambda x: IsDtypeKind.dtype(x))
+converter(str, IsDtypeKind, lambda x: IsDtypeKind.string(x))
+converter(
+    tuple, IsDtypeKind, lambda x: convert(x[0], IsDtypeKind) | convert(x[1:], IsDtypeKind) if x else IsDtypeKind.NULL
+)
+
+
+@array_api_ruleset.register
+def _isdtype(d: DType, k1: IsDtypeKind, k2: IsDtypeKind):
+    return [
+        rewrite(isdtype(DType.float32, IsDtypeKind.string("integral"))).to(FALSE),
+        rewrite(isdtype(DType.float64, IsDtypeKind.string("integral"))).to(FALSE),
+        rewrite(isdtype(DType.object, IsDtypeKind.string("integral"))).to(FALSE),
+        rewrite(isdtype(DType.int64, IsDtypeKind.string("integral"))).to(TRUE),
+        rewrite(isdtype(DType.int32, IsDtypeKind.string("integral"))).to(TRUE),
+        rewrite(isdtype(DType.float32, IsDtypeKind.string("real floating"))).to(TRUE),
+        rewrite(isdtype(DType.float64, IsDtypeKind.string("real floating"))).to(TRUE),
+        rewrite(isdtype(DType.object, IsDtypeKind.string("real floating"))).to(FALSE),
+        rewrite(isdtype(DType.int64, IsDtypeKind.string("real floating"))).to(FALSE),
+        rewrite(isdtype(DType.int32, IsDtypeKind.string("real floating"))).to(FALSE),
+        rewrite(isdtype(DType.float32, IsDtypeKind.string("complex floating"))).to(FALSE),
+        rewrite(isdtype(DType.float64, IsDtypeKind.string("complex floating"))).to(FALSE),
+        rewrite(isdtype(DType.object, IsDtypeKind.string("complex floating"))).to(FALSE),
+        rewrite(isdtype(DType.int64, IsDtypeKind.string("complex floating"))).to(FALSE),
+        rewrite(isdtype(DType.int32, IsDtypeKind.string("complex floating"))).to(FALSE),
+        rewrite(isdtype(d, IsDtypeKind.NULL)).to(FALSE),
+        rewrite(isdtype(d, IsDtypeKind.dtype(d))).to(TRUE),
+        rewrite(isdtype(d, k1 | k2)).to(isdtype(d, k1) | isdtype(d, k2)),
+        rewrite(k1 | IsDtypeKind.NULL).to(k1),
+    ]
+
+
 class Slice(Expr):
     def __init__(
         self,
@@ -1334,7 +1346,7 @@ def _unique(xs: TupleValue, a: NDArray, shape: TupleInt, copy: OptionalBool):
 
 @array_api_ruleset.register
 def _size(x: NDArray):
-    yield rewrite(x.size).to(x.shape.product())
+    yield rewrite(x.size).to(x.shape.fold(Int(1), Int.__mul__))
 
 
 @overload

egglog/exp/array_api_numba.py

@@ -31,7 +31,12 @@ 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)))
+    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)

egglog/exp/siu_examples.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+import egglog
+
+from .array_api import Int
+
+# https://github.com/sklam/pyasir/blob/c363ff4f8f91177700ad4108dd5042b9b97d8289/pyasir/tests/test_fib.py
+
+# In progress - should be able to re-create this
+# @df.func
+# def fib_ir(n: pyasir.Int64) -> pyasir.Int64:
+#     @df.switch(n <= 1)
+#     def swt(n):
+#         @df.case(1)
+#         def case0(n):
+#             return 1
+
+#         @df.case(0)
+#         def case1(n):
+#             return fib_ir(n - 1) + fib_ir(n - 2)
+
+#         yield case0
+#         yield case1
+
+#     r = swt(n)
+#     return r
+
+
+# With something like this:
+@egglog.function
+def fib(n: Int) -> Int:
+    return (n <= Int(1)).if_int(
+        Int(1),
+        fib(n - Int(1)) + fib(n - Int(2)),
+    )