Done?

Author: saulshanabrook

python/egglog/egraph_state.py

@@ -397,7 +397,7 @@ def callable_ref_to_egg(self, ref: CallableRef) -> tuple[str, bool]:  # noqa: C9
                     bindings.Relation(span(), egg_name, [self.type_ref_to_egg(a) for a in arg_types])
                 )
             case ConstantDecl(tp, _):
-                # Use constructor decleration instead of constant b/c constants cannot be extracted
+                # Use constructor declaration instead of constant b/c constants cannot be extracted
                 # https://github.com/egraphs-good/egglog/issues/334
                 is_function = self.__egg_decls__._classes[tp.ident].builtin
                 schema = bindings.Schema([], self.type_ref_to_egg(tp))

python/egglog/exp/array_api.py

@@ -836,7 +836,7 @@ def _tuple_int(
     yield rule(eq(ti).to(TupleInt(vs)), eq(ti).to(TupleInt(vs2)), vs != vs2).then(vs | vs2)
 
     yield rewrite(TupleInt.fn(i2, idx_fn).length(), subsume=False).to(i2)
-    yield rewrite(TupleInt.fn(i2, idx_fn)[i], subsume=False).to(idx_fn(check_index(i2, i)))
+    yield rewrite(TupleInt.fn(i2, idx_fn)[i], subsume=True).to(idx_fn(check_index(i2, i)))
 
     yield rewrite(TupleInt(vs).length()).to(Int(vs.length()))
     yield rewrite(TupleInt(vs)[Int(k)]).to(vs[k])
@@ -955,7 +955,7 @@ def _tuple_tuple_int(
     yield rule(eq(ti).to(TupleTupleInt(vs)), eq(ti).to(TupleTupleInt(vs2)), vs != vs2).then(vs | vs2)
 
     yield rewrite(TupleTupleInt.fn(i2, idx_fn).length(), subsume=False).to(i2)
-    yield rewrite(TupleTupleInt.fn(i2, idx_fn)[i], subsume=False).to(idx_fn(check_index(i2, i)))
+    yield rewrite(TupleTupleInt.fn(i2, idx_fn)[i], subsume=True).to(idx_fn(check_index(i2, i)))
 
     yield rewrite(TupleTupleInt(vs).length(), subsume=False).to(Int(vs.length()))
     yield rewrite(TupleTupleInt(vs)[Int(k)], subsume=False).to(vs[k])
@@ -1330,7 +1330,7 @@ def _tuple_value(
     yield rule(eq(ti).to(TupleValue(vs)), eq(ti).to(TupleValue(vs2)), vs != vs2).then(vs | vs2)
 
     yield rewrite(TupleValue.fn(i2, idx_fn).length(), subsume=False).to(i2)
-    yield rewrite(TupleValue.fn(i2, idx_fn)[i], subsume=False).to(idx_fn(check_index(i2, i)))
+    yield rewrite(TupleValue.fn(i2, idx_fn)[i], subsume=True).to(idx_fn(check_index(i2, i)))
 
     yield rewrite(TupleValue(vs).length(), subsume=False).to(Int(vs.length()))
     yield rewrite(TupleValue(vs)[Int(k)], subsume=False).to(vs[k], k >= 0, k < vs.length())
@@ -1961,7 +1961,7 @@ def _tuple_ndarray(
 ):
     yield rule(eq(ti).to(TupleNDArray(vs)), eq(ti).to(TupleNDArray(vs2)), vs != vs2).then(vs | vs2)
     yield rewrite(TupleNDArray.fn(i2, idx_fn).length(), subsume=False).to(i2)
-    yield rewrite(TupleNDArray.fn(i2, idx_fn)[i], subsume=False).to(idx_fn(check_index(i2, i)))
+    yield rewrite(TupleNDArray.fn(i2, idx_fn)[i], subsume=True).to(idx_fn(check_index(i2, i)))
 
     yield rewrite(TupleNDArray(vs).length(), subsume=False).to(Int(vs.length()))
     yield rewrite(TupleNDArray(vs)[Int(k)], subsume=False).to(vs[k], k >= 0, k < vs.length())
@@ -2148,26 +2148,33 @@ def _astype(x: NDArray, dtype: DType, i: i64):
     ]
 
 
-@function
+@function(unextractable=True, ruleset=array_api_ruleset)
 def unique_counts(x: NDArray) -> TupleNDArray:
     """
     Returns the unique elements of an input array x and the corresponding counts for each unique element in x.
 
 
     https://data-apis.org/array-api/2022.12/API_specification/generated/array_api.unique_counts.html
     """
+    return TupleNDArray((unique_counts_elements(x), unique_counts_counts(x)))
+
+
+@function
+def unique_counts_elements(x: NDArray) -> NDArray: ...
+
+
+@function
+def unique_counts_counts(x: NDArray) -> NDArray: ...
 
 
 @array_api_ruleset.register
 def _unique_counts(x: NDArray, c: NDArray, tv: TupleValue, v: Value, dtype: DType):
     return [
-        # rewrite(unique_counts(x).length()).to(Int(2)),
-        rewrite(unique_counts(x)).to(TupleNDArray.fn(2, unique_counts(x).__getitem__)),
         # Sum of all unique counts is the size of the array
-        rewrite(sum(unique_counts(x)[Int(1)])).to(NDArray(Value.from_int(x.size))),
+        rewrite(sum(unique_counts_counts(x))).to(NDArray(Value.from_int(x.size))),
         # Same but with astype in the middle
         # TODO: Replace
-        rewrite(sum(astype(unique_counts(x)[Int(1)], dtype))).to(astype(NDArray(Value.from_int(x.size)), dtype)),
+        rewrite(sum(astype(unique_counts_counts(x), dtype))).to(astype(NDArray(Value.from_int(x.size)), dtype)),
     ]
 
 
@@ -2194,22 +2201,25 @@ def _abs(f: Float):
     ]
 
 
-@function
+@function(ruleset=array_api_ruleset, unextractable=True)
 def unique_inverse(x: NDArray) -> TupleNDArray:
     """
     Returns the unique elements of an input array x and the indices from the set of unique elements that reconstruct x.
 
     https://data-apis.org/array-api/2022.12/API_specification/generated/array_api.unique_inverse.html
     """
+    return TupleNDArray((unique_values(x), unique_inverse_inverse_indices(x)))
+
+
+@function
+def unique_inverse_inverse_indices(x: NDArray) -> NDArray: ...
 
 
 @array_api_ruleset.register
 def _unique_inverse(x: NDArray, i: Int):
     return [
-        # rewrite(unique_inverse(x).length()).to(Int(2)),
-        rewrite(unique_inverse(x)).to(TupleNDArray.fn(2, unique_inverse(x).__getitem__)),
         # Shape of unique_inverse first element is same as shape of unique_values
-        rewrite(unique_inverse(x)[Int(0)]).to(unique_values(x)),
+        rewrite(unique_values(x)[Int(0)]).to(unique_values(x)),
     ]
 
 
@@ -2319,19 +2329,19 @@ def cross(a: NDArrayLike, b: NDArrayLike) -> NDArray:
 linalg = sys.modules[__name__]
 
 
-@function
-def svd(x: NDArray, full_matrices: Boolean = TRUE) -> TupleNDArray:
+def svd(x: NDArray, full_matrices: Boolean = TRUE) -> tuple[NDArray, NDArray, NDArray]:
     """
     https://data-apis.org/array-api/2022.12/extensions/generated/array_api.linalg.svd.html
     """
+    res = svd_(x, full_matrices)
+    return (res[0], res[1], res[2])
 
 
-@array_api_ruleset.register
-def _linalg(x: NDArray, full_matrices: Boolean):
-    return [
-        # rewrite(svd(x, full_matrices).length()).to(Int(3)),
-        rewrite(svd(x, full_matrices)).to(TupleNDArray.fn(3, svd(x, full_matrices).__getitem__)),
-    ]
+@function
+def svd_(x: NDArray, full_matrices: Boolean = TRUE) -> TupleNDArray:
+    """
+    https://data-apis.org/array-api/2022.12/extensions/generated/array_api.linalg.svd.html
+    """
 
 
 @function(ruleset=array_api_ruleset, unextractable=True)
@@ -2427,7 +2437,7 @@ def _interval_analaysis(
         # rule(eq(y).to(any(x)), ndarray_all_false(x)).then(union(y).with_(NDArray(Value.bool(FALSE)))),
         # Indexing into unique counts counts are all positive
         rule(
-            eq(v).to(unique_counts(x)[Int(1)].index(idx)),
+            eq(v).to(unique_counts_counts(x).index(idx)),
         ).then(greater_zero(v)),
         # Min value preserved over astype
         rule(
@@ -2710,15 +2720,20 @@ def try_evaling(expr: ExprWithValue[T_co]) -> T_co:
     egraph = _get_current_egraph()
     egraph.register(expr)  # type: ignore[arg-type]
     egraph.run(array_api_schedule)
-    # run on another e-graph to get around bug
-    # https://github.com/egraphs-good/egglog/issues/801
-    # return egraph.extract(expr).value  # type: ignore[call-overload]
     # egraph.display(n_inline_leaves=2, split_primitive_outputs=True, split_functions=[Int])
     extracted_expr = egraph.extract(expr)  # type: ignore[call-overload]
+    with contextlib.suppress(ExprValueError):
+        extracted_expr.value
+    # run on another e-graph to get around bug
+    # https://github.com/egraphs-good/egglog/issues/801
+
     new_egraph = EGraph()
     new_egraph.register(extracted_expr)
     new_egraph.run(array_api_schedule)
     return new_egraph.extract(extracted_expr).value
+    # except EggSmolError as e:
+
+    #     raise e
 
     # try:
     #     return egraph.extract(prim_expr).value  # type: ignore[attr-defined]

python/egglog/exp/array_api_jit.py

@@ -4,7 +4,7 @@
 
 import numpy as np
 
-from egglog import EGraph, greedy_dag_cost_model
+from egglog import EGraph, bindings, greedy_dag_cost_model
 from egglog.exp.array_api import NDArray, set_array_api_egraph
 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
@@ -24,14 +24,21 @@ def jit(
     Jit compiles a function
     """
     egraph, res, res_optimized, program = function_to_program(fn, save_egglog_string=False)
+    egraph = EGraph()
     if handle_expr:
         handle_expr(res)
     if handle_optimized_expr:
         handle_optimized_expr(res_optimized)
     fn_program = EvalProgram(program, {"np": np})
     egraph.register(fn_program)
+
     egraph.run(array_api_program_gen_schedule)
-    return cast("X", egraph.extract(fn_program.as_py_object).value)
+
+    try:
+        return cast("X", egraph.extract(fn_program.as_py_object).value)
+    except bindings.EggSmolError as e:
+        e.add_note(f"Failed to get py object from {egraph.extract(fn_program)}")
+        raise
 
 
 def function_to_program(fn: Callable, save_egglog_string: bool) -> tuple[EGraph, NDArray, NDArray, Program]:

python/egglog/exp/array_api_numba.py

@@ -49,7 +49,7 @@ def count_values(x: NDArrayLike, values: TupleValueLike) -> TupleValue:
 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(
+        rewrite(unique_counts_counts(x), subsume=True).to(
             NDArray.from_tuple_value(count_values(x, unique_values(x).to_tuple_values()))
         ),
     ]
@@ -60,5 +60,5 @@ def _unique_counts(x: NDArray, c: NDArray, tv: TupleValue, v: Value):
 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)[1] == i, subsume=True).to(x == unique_values(x).index((i,))),
+        rewrite(unique_inverse_inverse_indices(x) == i, subsume=True).to(x == unique_values(x).index((i,))),
     ]

python/egglog/exp/array_api_program_gen.py

@@ -35,7 +35,7 @@ def int_program(x: Int) -> Program: ...
 
 
 @array_api_program_gen_ruleset.register
-def _int_program(i64_: i64, i: Int, j: Int, s: String):
+def _int_program(i64_: i64, i: Int, j: Int, s: String, b: Boolean, ti: Callable[[], Int], ti1: Callable[[], Int]):
     yield rewrite(int_program(Int.var(s))).to(Program(s, True))
     yield rewrite(int_program(Int(i64_))).to(Program(i64_.to_string()))
     yield rewrite(int_program(~i)).to(Program("~") + int_program(i))
@@ -57,6 +57,15 @@ def _int_program(i64_: i64, i: Int, j: Int, s: String):
     yield rewrite(int_program(i >> j)).to(Program("(") + int_program(i) + " >> " + int_program(j) + ")")
     yield rewrite(int_program(i // j)).to(Program("(") + int_program(i) + " // " + int_program(j) + ")")
 
+    assigned = int_program(j).assign()
+    yield rewrite(int_program(check_index(i, j)), subsume=True).to(
+        assigned.statement(Program("assert ") + assigned + " < " + int_program(i))
+    )
+
+    yield rewrite(int_program(Int.if_(b, ti, ti1))).to(
+        int_program(ti()) + " if " + bool_program(b) + " else " + int_program(ti1())
+    )
+
 
 @function
 def program_if(b: BooleanLike, t: Callable[[], Program], f: Callable[[], Program]) -> Program: ...
@@ -168,7 +177,10 @@ def _value_program(i: Int, b: Boolean, f: Float, x: NDArray, v1: Value, v2: Valu
     yield rewrite(value_program(v1 * v2)).to(Program("(") + value_program(v1) + " * " + value_program(v2) + ")")
     yield rewrite(bool_program(v1.to_bool)).to(value_program(v1))
     yield rewrite(int_program(v1.to_int)).to(value_program(v1))
-    yield rewrite(value_program(xs.index(ti))).to((ndarray_program(xs) + "[" + tuple_int_program(ti) + "]").assign())
+    yield rewrite(value_program(xs.index(ti))).to(
+        (ndarray_program(xs) + "[" + tuple_int_program(ti) + "]").assign(), ne(ti).to(TupleInt(()))
+    )
+    yield rewrite(value_program(xs.index(TupleInt(())))).to(ndarray_program(xs))
     yield rewrite(value_program(v1.sqrt())).to(Program("np.sqrt(") + value_program(v1) + ")")
     yield rewrite(value_program(v1.real())).to(Program("np.real(") + value_program(v1) + ")")
     yield rewrite(value_program(v1.conj())).to(Program("np.conj(") + value_program(v1) + ")")
@@ -516,8 +528,8 @@ def bin_op(res: NDArray, op: str) -> Command:
         optional_int_or_tuple_ != OptionalIntOrTuple.none,
     )
     # svd
-    yield rewrite(tuple_ndarray_program(svd(x))).to((Program("np.linalg.svd(") + ndarray_program(x) + ")").assign())
-    yield rewrite(tuple_ndarray_program(svd(x, FALSE))).to(
+    yield rewrite(tuple_ndarray_program(svd_(x))).to((Program("np.linalg.svd(") + ndarray_program(x) + ")").assign())
+    yield rewrite(tuple_ndarray_program(svd_(x, FALSE))).to(
         (Program("np.linalg.svd(") + ndarray_program(x) + ", full_matrices=False)").assign()
     )
     # sqrt
@@ -564,4 +576,5 @@ def _vec_recursive_value_program(v: Value, vv: Vec[RecursiveValue]):
     yield rewrite(vec_recursive_value_program(Vec[RecursiveValue].empty())).to(Program(""))
     yield rewrite(vec_recursive_value_program(vv)).to(
         recursive_value_program(vv[0]) + ", " + vec_recursive_value_program(vv.remove(0)),
+        vv.length() > 0,
     )

python/tests/__snapshots__/test_array_api/test_jit[lda][code].py

@@ -0,0 +1,64 @@
+def __fn(X, y):
+    assert X.dtype == np.dtype(np.float64)
+    assert X.shape == (150, 4, )
+    assert np.all(np.isfinite(X))
+    assert y.dtype == np.dtype(np.int64)
+    assert y.shape == (150, )
+    assert set(np.unique(y)) == set((0, 1, 2, ))
+    _0 = y == np.array(0)
+    _1 = np.sum(_0)
+    _2 = y == np.array(1)
+    _3 = np.sum(_2)
+    _4 = y == np.array(2)
+    _5 = np.sum(_4)
+    _6 = np.array((_1, _3, _5, )).astype(np.dtype(np.float64))
+    _7 = _6 / np.array(150)
+    _8 = np.zeros((3, 4, ), dtype=np.dtype(np.float64))
+    _9 = np.sum(X[_0], axis=0)
+    _10 = _9 / np.array(X[_0].shape[0])
+    _8[0, :,] = _10
+    _11 = np.sum(X[_2], axis=0)
+    _12 = _11 / np.array(X[_2].shape[0])
+    _8[1, :,] = _12
+    _13 = np.sum(X[_4], axis=0)
+    _14 = _13 / np.array(X[_4].shape[0])
+    _8[2, :,] = _14
+    _15 = _7 @ _8
+    _16 = X - _15
+    _17 = np.sqrt(np.asarray(np.array(float(1 / 147)), np.dtype(np.float64)))
+    _18 = X[_0] - _8[0, :,]
+    _19 = X[_2] - _8[1, :,]
+    _20 = X[_4] - _8[2, :,]
+    _21 = np.concatenate((_18, _19, _20, ), axis=0)
+    _22 = np.sum(_21, axis=0)
+    _23 = _22 / np.array(_21.shape[0])
+    _24 = np.expand_dims(_23, 0)
+    _25 = _21 - _24
+    _26 = np.square(_25)
+    _27 = np.sum(_26, axis=0)
+    _28 = _27 / np.array(_26.shape[0])
+    _29 = np.sqrt(_28)
+    _30 = _29 == np.array(0)
+    _29[_30] = np.array((150 / 150))
+    _31 = _21 / _29
+    _32 = _17 * _31
+    _33 = np.linalg.svd(_32, full_matrices=False)
+    _34 = _33[1] > np.array(0.0001)
+    _35 = _34.astype(np.dtype(np.int32))
+    _36 = np.sum(_35)
+    _37 = _33[2][:_36, :,] / _29
+    _38 = _37.T / _33[1][:_36]
+    _39 = np.array(150) * _7
+    _40 = _39 * np.array(float(1 / 2))
+    _41 = np.sqrt(_40)
+    _42 = _8 - _15
+    _43 = _41 * _42.T
+    _44 = _43.T @ _38
+    _45 = np.linalg.svd(_44, full_matrices=False)
+    _46 = np.array(0.0001) * _45[1][0]
+    _47 = _45[1] > _46
+    _48 = _47.astype(np.dtype(np.int32))
+    _49 = np.sum(_48)
+    _50 = _38 @ _45[2].T[:, :_49,]
+    _51 = _16 @ _50
+    return _51[:, :2,]