egglog 10.0.1__cp313-cp313-win_amd64.whl → 11.0.0__cp313-cp313-win_amd64.whl

egglog/exp/array_api.py

@@ -69,6 +69,7 @@ import numpy as np
 
 from egglog import *
 from egglog.runtime import RuntimeExpr
+from egglog.version_compat import add_note
 
 from .program_gen import *
 
@@ -153,6 +154,18 @@ class Int(Expr, ruleset=array_api_ruleset):
     def __eq__(self, other: IntLike) -> Boolean:  # type: ignore[override]
         ...
 
+    # add a hash so that this test can pass
+    # https://github.com/scikit-learn/scikit-learn/blob/6fd23fca53845b32b249f2b36051c081b65e2fab/sklearn/utils/validation.py#L486-L487
+    @method(preserve=True)
+    def __hash__(self) -> int:
+        egraph = _get_current_egraph()
+        egraph.register(self)
+        egraph.run(array_api_schedule)
+        simplified = egraph.extract(self)
+        return hash(cast("RuntimeExpr", simplified).__egg_typed_expr__)
+
+    def __round__(self, ndigits: OptionalIntLike = None) -> Int: ...
+
     # TODO: Fix this?
     # Make != always return a Bool, so that numpy.unique works on a tuple of ints
     # In _unique1d
@@ -279,6 +292,8 @@ def _int(i: i64, j: i64, r: Boolean, o: Int, b: Int):
     yield rewrite(Int.if_(TRUE, o, b), subsume=True).to(o)
     yield rewrite(Int.if_(FALSE, o, b), subsume=True).to(b)
 
+    yield rewrite(o.__round__(OptionalInt.none)).to(o)
+
     # Never cannot be equal to anything real
     yield rule(eq(Int.NEVER).to(Int(i))).then(panic("Int.NEVER cannot be equal to any real int"))
 
@@ -353,8 +368,14 @@ class Float(Expr, ruleset=array_api_ruleset):
     def __sub__(self, other: FloatLike) -> Float: ...
 
     def __pow__(self, other: FloatLike) -> Float: ...
+    def __round__(self, ndigits: OptionalIntLike = None) -> Float: ...
 
     def __eq__(self, other: FloatLike) -> Boolean: ...  # type: ignore[override]
+    def __ne__(self, other: FloatLike) -> Boolean: ...  # type: ignore[override]
+    def __lt__(self, other: FloatLike) -> Boolean: ...
+    def __le__(self, other: FloatLike) -> Boolean: ...
+    def __gt__(self, other: FloatLike) -> Boolean: ...
+    def __ge__(self, other: FloatLike) -> Boolean: ...
 
 
 converter(float, Float, lambda x: Float(x))
@@ -365,9 +386,10 @@ FloatLike: TypeAlias = Float | float | IntLike
 
 
 @array_api_ruleset.register
-def _float(fl: Float, f: f64, f2: f64, i: i64, r: BigRat, r1: BigRat):
+def _float(fl: Float, f: f64, f2: f64, i: i64, r: BigRat, r1: BigRat, i_: Int):
     return [
         rule(eq(fl).to(Float(f))).then(set_(fl.to_f64).to(f)),
+        rewrite(Float.from_int(Int(i))).to(Float(f64.from_i64(i))),
         rewrite(Float(f).abs()).to(Float(f), f >= 0.0),
         rewrite(Float(f).abs()).to(Float(-f), f < 0.0),
         # Convert from float to rationl, if its a whole number i.e. can be converted to int
@@ -382,11 +404,22 @@ def _float(fl: Float, f: f64, f2: f64, i: i64, r: BigRat, r1: BigRat):
         rewrite(Float.rational(r) - Float.rational(r1)).to(Float.rational(r - r1)),
         rewrite(Float.rational(r) * Float.rational(r1)).to(Float.rational(r * r1)),
         rewrite(Float(f) ** Float(f2)).to(Float(f**f2)),
-        # ==
+        # comparisons
         rewrite(Float(f) == Float(f)).to(TRUE),
         rewrite(Float(f) == Float(f2)).to(FALSE, ne(f).to(f2)),
+        rewrite(Float(f) != Float(f2)).to(TRUE, f != f2),
+        rewrite(Float(f) != Float(f)).to(FALSE),
+        rewrite(Float(f) >= Float(f2)).to(TRUE, f >= f2),
+        rewrite(Float(f) >= Float(f2)).to(FALSE, f < f2),
+        rewrite(Float(f) <= Float(f2)).to(TRUE, f <= f2),
+        rewrite(Float(f) <= Float(f2)).to(FALSE, f > f2),
+        rewrite(Float(f) > Float(f2)).to(TRUE, f > f2),
+        rewrite(Float(f) > Float(f2)).to(FALSE, f <= f2),
+        rewrite(Float(f) < Float(f2)).to(TRUE, f < f2),
         rewrite(Float.rational(r) == Float.rational(r)).to(TRUE),
         rewrite(Float.rational(r) == Float.rational(r1)).to(FALSE, ne(r).to(r1)),
+        # round
+        rewrite(Float.rational(r).__round__()).to(Float.rational(r.round())),
     ]
 
 
@@ -670,6 +703,8 @@ class OptionalInt(Expr, ruleset=array_api_ruleset):
     def some(cls, value: Int) -> OptionalInt: ...
 
 
+OptionalIntLike: TypeAlias = OptionalInt | IntLike | None
+
 converter(type(None), OptionalInt, lambda _: OptionalInt.none)
 converter(Int, OptionalInt, OptionalInt.some)
 
@@ -1198,13 +1233,13 @@ class NDArray(Expr, ruleset=array_api_ruleset):
 
 NDArrayLike: TypeAlias = NDArray | ValueLike | TupleValueLike
 
-converter(NDArray, IndexKey, IndexKey.ndarray)
-converter(Value, NDArray, NDArray.scalar)
+converter(NDArray, IndexKey, lambda v: IndexKey.ndarray(v))
+converter(Value, NDArray, lambda v: NDArray.scalar(v))
 # Need this if we want to use ints in slices of arrays coming from 1d arrays, but make it more expensive
 # to prefer upcasting in the other direction when we can, which is safer at runtime
 converter(NDArray, Value, lambda n: n.to_value(), 100)
-converter(TupleValue, NDArray, NDArray.vector)
-converter(TupleInt, TupleValue, TupleValue.from_tuple_int)
+converter(TupleValue, NDArray, lambda v: NDArray.vector(v))
+converter(TupleInt, TupleValue, lambda v: TupleValue.from_tuple_int(v))
 
 
 @array_api_ruleset.register
@@ -1383,8 +1418,8 @@ class IntOrTuple(Expr, ruleset=array_api_ruleset):
     def tuple(cls, value: TupleIntLike) -> IntOrTuple: ...
 
 
-converter(Int, IntOrTuple, IntOrTuple.int)
-converter(TupleInt, IntOrTuple, IntOrTuple.tuple)
+converter(Int, IntOrTuple, lambda v: IntOrTuple.int(v))
+converter(TupleInt, IntOrTuple, lambda v: IntOrTuple.tuple(v))
 
 
 class OptionalIntOrTuple(Expr, ruleset=array_api_ruleset):
@@ -1395,7 +1430,7 @@ class OptionalIntOrTuple(Expr, ruleset=array_api_ruleset):
 
 
 converter(type(None), OptionalIntOrTuple, lambda _: OptionalIntOrTuple.none)
-converter(IntOrTuple, OptionalIntOrTuple, OptionalIntOrTuple.some)
+converter(IntOrTuple, OptionalIntOrTuple, lambda v: OptionalIntOrTuple.some(v))
 
 
 @function
@@ -1980,6 +2015,5 @@ def try_evaling(egraph: EGraph, schedule: Schedule, expr: Expr, prim_expr: Built
             extracted = egraph.extract(prim_expr)
         except BaseException as e:
             # egraph.display(n_inline_leaves=1, split_primitive_outputs=True)
-            e.add_note(f"Cannot evaluate {egraph.extract(expr)}")
-            raise
-    return extracted.eval()  # type: ignore[attr-defined]
+            raise add_note(f"Cannot evaluate {egraph.extract(expr)}", e)  # noqa: B904
+    return extracted.value  # type: ignore[attr-defined]

egglog/exp/array_api_jit.py

@@ -14,16 +14,22 @@ from .program_gen import Program
 X = TypeVar("X", bound=Callable)
 
 
-def jit(fn: X) -> X:
+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})
-    fn = cast("X", try_evaling(egraph, array_api_program_gen_schedule, fn_program, fn_program.as_py_object))
-    fn.initial_expr = res  # type: ignore[attr-defined]
-    fn.expr = res_optimized  # type: ignore[attr-defined]
-    return fn
+    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]:

egglog/exp/array_api_program_gen.py

@@ -505,6 +505,6 @@ def _ndarray_program(
     yield rewrite(ndarray_program(abs(x))).to((Program("np.abs(") + ndarray_program(x) + ")").assign())
 
     # asarray
-    yield rewrite(ndarray_program(asarray(x, odtype))).to(
+    yield rewrite(ndarray_program(asarray(x, odtype, OptionalBool.none, optional_device_))).to(
         Program("np.asarray(") + ndarray_program(x) + ", " + optional_dtype_program(odtype) + ")"
     )

egglog/exp/program_gen.py

@@ -5,12 +5,10 @@ Builds up imperative string expressions from a functional expression.
 
 from __future__ import annotations
 
-from typing import TypeAlias, Union
+from typing import TypeAlias
 
 from egglog import *
 
-ProgramLike: TypeAlias = Union["Program", StringLike]
-
 
 class Program(Expr):
     """
@@ -97,6 +95,9 @@ class Program(Expr):
         """
 
 
+ProgramLike: TypeAlias = Program | StringLike
+
+
 converter(String, Program, Program)
 
 

egglog/pretty.py

@@ -107,7 +107,7 @@ def pretty_callable_ref(
     """
     # Pass in three dummy args, which are the max used for any operation that
     # is not a generic function call
-    args: list[ExprDecl] = [VarDecl(ARG_STR, False)] * 3
+    args: list[ExprDecl] = [UnboundVarDecl(ARG_STR)] * 3
     if first_arg:
         args.insert(0, first_arg)
     context = PrettyContext(decls, defaultdict(lambda: 0))
@@ -166,7 +166,7 @@ class TraverseContext:
                 self(d.expr)
             case ChangeDecl(_, d, _) | SaturateDecl(d) | RepeatDecl(d, _) | ActionCommandDecl(d):
                 self(d)
-            case PanicDecl(_) | VarDecl(_) | LitDecl(_) | PyObjectDecl(_):
+            case PanicDecl(_) | UnboundVarDecl(_) | LetRefDecl(_) | LitDecl(_) | PyObjectDecl(_):
                 pass
             case SequenceDecl(decls) | RulesetDecl(decls):
                 for de in decls:
@@ -233,6 +233,10 @@ class PrettyContext:
         return expr
 
     def uncached(self, decl: AllDecls, *, unwrap_lit: bool, parens: bool, ruleset_name: str | None) -> tuple[str, str]:  # noqa: C901, PLR0911, PLR0912
+        """
+        Returns a tuple of a string value of the decleration and the "type" to use when create a memoized cached version
+        for de-duplication.
+        """
         match decl:
             case LitDecl(value):
                 match value:
@@ -247,7 +251,7 @@ class PrettyContext:
                     case str(s):
                         return repr(s) if unwrap_lit else f"String({s!r})", "String"
                 assert_never(value)
-            case VarDecl(name):
+            case UnboundVarDecl(name) | LetRefDecl(name):
                 return name, name
             case CallDecl(_, _, _):
                 return self._call(decl, parens)
@@ -357,7 +361,7 @@ class PrettyContext:
             has_multiple_parents = self.parents[first_arg] > 1
             self.names[decl] = expr_name = self._name_expr(tp_name, expr_str, copy_identifier=has_multiple_parents)
             # Set the first arg to be the name of the mutated arg and return the name
-            args[0] = VarDecl(expr_name, True)
+            args[0] = LetRefDecl(expr_name)
         else:
             expr_name = None
         res = self._call_inner(ref, args, decl.bound_tp_params, parens)
@@ -390,6 +394,7 @@ class PrettyContext:
                 return f"{tp_ref}.{method_name}", args
             case MethodRef(_class_name, method_name):
                 slf, *args = args
+                non_str_slf = slf
                 slf = self(slf, parens=True)
                 match method_name:
                     case _ if method_name in UNARY_METHODS:
@@ -406,6 +411,8 @@ class PrettyContext:
                         return f"del {slf}[{self(args[0], unwrap_lit=True)}]"
                     case "__setitem__":
                         return f"{slf}[{self(args[0], unwrap_lit=True)}] = {self(args[1], unwrap_lit=True)}"
+                    case "__round__":
+                        return "round", [non_str_slf, *args]
                     case _:
                         return f"{slf}.{method_name}", args
             case ConstantRef(name):
@@ -487,24 +494,3 @@ class PrettyContext:
         if arg_names:
             prefix += f" {', '.join(self(a.expr) for a in arg_names)}"
         return f"{prefix}: {self(res.expr)}"
-
-
-def _plot_line_length(expr: object):  # pragma: no cover
-    """
-    Plots the number of line lengths based on different max lengths
-    """
-    global MAX_LINE_LENGTH, LINE_DIFFERENCE
-    import altair as alt
-    import pandas as pd
-
-    sizes = []
-    for line_length in range(40, 180, 10):
-        MAX_LINE_LENGTH = line_length
-        for diff in range(0, 40, 5):
-            LINE_DIFFERENCE = diff
-            new_l = len(str(expr).split())
-            sizes.append((line_length, diff, new_l))
-
-    df = pd.DataFrame(sizes, columns=["MAX_LINE_LENGTH", "LENGTH_DIFFERENCE", "n"])
-
-    return alt.Chart(df).mark_rect().encode(x="MAX_LINE_LENGTH:O", y="LENGTH_DIFFERENCE:O", color="n:Q")