Start adding support for getting actual value of indexed array.

Author: saulshanabrook

python/egglog/exp/array_api.py

@@ -303,6 +303,7 @@ def __len__(self) -> int:
     def __iter__(self) -> Iterator[Int]:
         return iter(self[i] for i in range(len(self)))
 
+    # TODO: Rename to reduce to match Python? And re-order?
     def fold(self, init: Int, f: Callable[[Int, Int], Int]) -> Int: ...
 
     def fold_boolean(self, init: Boolean, f: Callable[[Boolean, Int], Boolean]) -> Boolean: ...
@@ -861,7 +862,8 @@ def __rxor__(self, other: NDArray) -> NDArray: ...
     def __ror__(self, other: NDArray) -> NDArray: ...
 
     @classmethod
-    def scalar(cls, value: Value) -> NDArray: ...
+    def scalar(cls, value: Value) -> NDArray:
+        return NDArray(TupleInt.EMPTY, value.dtype, lambda _: value)
 
     def to_value(self) -> Value: ...
 
@@ -900,10 +902,12 @@ def _ndarray(
     shape: TupleInt,
     dtype: DType,
     idx_fn: Callable[[TupleInt], Value],
+    idx: TupleInt,
 ):
     return [
         rewrite(NDArray(shape, dtype, idx_fn).shape).to(shape),
         rewrite(NDArray(shape, dtype, idx_fn).dtype).to(dtype),
+        rewrite(NDArray(shape, dtype, idx_fn).index(idx), subsume=True).to(idx_fn(idx)),
         rewrite(x.ndim).to(x.shape.length()),
         # rewrite(NDArray.scalar(Value.bool(b)).to_bool()).to(b),
         # Converting to a value requires a scalar bool value

python/egglog/exp/array_api_loopnest.py

@@ -69,12 +69,63 @@ def tuple_tuple_int_reduce_ndarray(
 ) -> NDArray: ...
 
 
+@array_api_ruleset.register
+def _tuple_tuple_int_reduce_ndarray(
+    idx_fn: Callable[[Int], TupleInt], f: Callable[[NDArray, TupleInt], NDArray], i: NDArray, k: i64
+):
+    yield rewrite(tuple_tuple_int_reduce_ndarray(TupleTupleInt(0, idx_fn), f, i)).to(i)
+    yield rewrite(tuple_tuple_int_reduce_ndarray(TupleTupleInt(Int(k), idx_fn), f, i), subsume=True).to(
+        f(
+            tuple_tuple_int_reduce_ndarray(TupleTupleInt(k - 1, lambda i: idx_fn(i + 1)), f, i),
+            idx_fn(Int(0)),
+        ),
+        ne(k).to(i64(0)),
+    )
+
+
 @function
 def tuple_int_map_tuple_int(xs: TupleInt, fn: Callable[[Int], TupleInt]) -> TupleTupleInt: ...
 
 
+@array_api_ruleset.register
+def _tuple_int_map_tuple_int(length: Int, fn: Callable[[Int], TupleInt], idx_fn: Callable[[Int], Int]):
+    yield rewrite(
+        tuple_int_map_tuple_int(
+            TupleInt(length, idx_fn),
+            fn,
+        )
+    ).to(TupleTupleInt(length, lambda i: fn(idx_fn(i))))
+
+
 @function
-def tuple_tuple_int_product(xs: TupleTupleInt) -> TupleTupleInt: ...
+def tuple_tuple_int_map_int(xs: TupleTupleInt, fn: Callable[[TupleInt], Int]) -> TupleInt: ...
+
+
+@array_api_ruleset.register
+def _tuple_tuple_int_map_int(length: Int, fn: Callable[[TupleInt], Int], idx_fn: Callable[[Int], TupleInt]):
+    yield rewrite(
+        tuple_tuple_int_map_int(
+            TupleTupleInt(length, idx_fn),
+            fn,
+        )
+    ).to(TupleInt(length, lambda i: fn(idx_fn(i))))
+
+
+@function
+def tuple_tuple_int_product_index(xs: TupleTupleInt, i: Int) -> TupleInt: ...
+
+
+@function(subsume=True, ruleset=array_api_ruleset)
+def tuple_tuple_int_product(xs: TupleTupleInt) -> TupleTupleInt:
+    """
+    Cartesian product of inputs
+
+    https://docs.python.org/3/library/itertools.html#itertools.product
+    """
+    # length is product of lengths
+    length = tuple_tuple_int_map_int(xs, lambda x: x.length()).fold(Int(1), lambda x, y: x * y)
+
+    return TupleTupleInt(length, partial(tuple_tuple_int_product_index, xs))
 
 
 @array_api_ruleset.register
@@ -107,3 +158,5 @@ def _loopnest_api_ruleset(
     yield rewrite(lna.indices, subsume=True).to(
         tuple_tuple_int_product(tuple_int_map_tuple_int(lna.get_dims(), TupleInt.range))
     )
+    # unwrap
+    yield rewrite(OptionalLoopNestAPI(lna).unwrap()).to(lna)

python/tests/test_array_api.py

@@ -94,6 +94,19 @@ def linalg_norm(X: NDArray, axis: TupleIntLike) -> NDArray:
     )
 
 
+X = NDArray.var("X")
+assume_shape(X, (3, 2, 3, 4))
+val = linalg_norm(X, (0, 1))
+i = constant("i", Int)
+j = constant("j", Int)
+idxed = val.index((i, j))
+
+egraph = EGraph()
+egraph.register(idxed)
+egraph.run(array_api_schedule)
+print(egraph.extract(idxed))
+
+
 class TestLoopNest:
     def test_shape(self):
         X = NDArray.var("X")
@@ -222,18 +235,18 @@ def test_execution(self, fn, benchmark):
 
 # if calling as script, print out egglog source for test
 # similar to jit, but don't include pyobject parts so it works in vanilla egglog
-if __name__ == "__main__":
-    print("Generating egglog source for test")
-    egraph = EGraph(save_egglog_string=True)
-    X_ = NDArray.var("X")
-    y_ = NDArray.var("y")
-    with egraph:
-        expr = lda(X_, y_)
-    optimized_expr = egraph.simplify(expr, array_api_numba_schedule)
-    fn_program = ndarray_function_two_program(optimized_expr, X_, y_)
-    egraph.register(fn_program.compile())
-    egraph.run(array_api_program_gen_ruleset.saturate() + program_gen_ruleset.saturate())
-    egraph.extract(fn_program.statements)
-    name = "python.egg"
-    print("Saving to", name)
-    Path(name).write_text(egraph.as_egglog_string)
+# if __name__ == "__main__":
+#     print("Generating egglog source for test")
+#     egraph = EGraph(save_egglog_string=True)
+#     X_ = NDArray.var("X")
+#     y_ = NDArray.var("y")
+#     with egraph:
+#         expr = lda(X_, y_)
+#     optimized_expr = egraph.simplify(expr, array_api_numba_schedule)
+#     fn_program = ndarray_function_two_program(optimized_expr, X_, y_)
+#     egraph.register(fn_program.compile())
+#     egraph.run(array_api_program_gen_ruleset.saturate() + program_gen_ruleset.saturate())
+#     egraph.extract(fn_program.statements)
+#     name = "python.egg"
+#     print("Saving to", name)
+#     Path(name).write_text(egraph.as_egglog_string)