egglog 6.0.1__cp310-none-win_amd64.whl → 6.1.0__cp310-none-win_amd64.whl

egglog/bindings.pyi

@@ -150,6 +150,20 @@ class Fact:
 
 _Fact: TypeAlias = Fact | Eq
 
+##
+# Change
+##
+
+@final
+class Delete:
+    def __init__(self) -> None: ...
+
+@final
+class Subsume:
+    def __init__(self) -> None: ...
+
+_Change: TypeAlias = Delete | Subsume
+
 ##
 # Actions
 ##
@@ -168,10 +182,11 @@ class Set:
     rhs: _Expr
 
 @final
-class Delete:
+class Change:
+    change: _Change
     sym: str
     args: list[_Expr]
-    def __init__(self, sym: str, args: list[_Expr]) -> None: ...
+    def __init__(self, change: _Change, sym: str, args: list[_Expr]) -> None: ...
 
 @final
 class Union:
@@ -195,7 +210,7 @@ class Extract:
     expr: _Expr
     variants: _Expr
 
-_Action: TypeAlias = Let | Set | Delete | Union | Panic | Expr_ | Extract
+_Action: TypeAlias = Let | Set | Change | Union | Panic | Expr_ | Extract
 
 ##
 # Other Structs
@@ -210,6 +225,7 @@ class FunctionDecl:
     merge_action: list[_Action]
     cost: int | None
     unextractable: bool
+    ignore_viz: bool
 
     def __init__(
         self,
@@ -220,6 +236,7 @@ class FunctionDecl:
         merge_action: list[_Action] = [],  # noqa: B006
         cost: int | None = None,
         unextractable: bool = False,
+        ignore_viz: bool = False,
     ) -> None: ...
 
 @final
@@ -374,7 +391,8 @@ class RewriteCommand:
     # TODO: Rename to ruleset
     name: str
     rewrite: Rewrite
-    def __init__(self, name: str, rewrite: Rewrite) -> None: ...
+    subsume: bool
+    def __init__(self, name: str, rewrite: Rewrite, subsume: bool) -> None: ...
 
 @final
 class BiRewriteCommand:

egglog/egraph.py

@@ -1465,6 +1465,7 @@ class Rewrite(Command):
     _lhs: RuntimeExpr
     _rhs: RuntimeExpr
     _conditions: tuple[Fact, ...]
+    _subsume: bool
     _fn_name: ClassVar[str] = "rewrite"
 
     def __str__(self) -> str:
@@ -1473,7 +1474,7 @@ class Rewrite(Command):
 
     def _to_egg_command(self, default_ruleset_name: str) -> bindings._Command:
         return bindings.RewriteCommand(
-            self.ruleset.egg_name if self.ruleset else default_ruleset_name, self._to_egg_rewrite()
+            self.ruleset.egg_name if self.ruleset else default_ruleset_name, self._to_egg_rewrite(), self._subsume
         )
 
     def _to_egg_rewrite(self) -> bindings.Rewrite:
@@ -1488,7 +1489,7 @@ class Rewrite(Command):
         return Declarations.create(self._lhs, self._rhs, *self._conditions)
 
     def with_ruleset(self, ruleset: Ruleset) -> Rewrite:
-        return Rewrite(ruleset, self._lhs, self._rhs, self._conditions)
+        return Rewrite(ruleset, self._lhs, self._rhs, self._conditions, self._subsume)
 
 
 @dataclass
@@ -1652,21 +1653,23 @@ class ExprAction(Action):
 
 
 @dataclass
-class Delete(Action):
+class Change(Action):
     """
-    Remove a function call from an EGraph.
+    Change a function call in an EGraph.
     """
 
+    change: Literal["delete", "subsume"]
     _call: RuntimeExpr
 
     def __str__(self) -> str:
-        return f"delete({self._call})"
+        return f"{self.change}({self._call})"
 
-    def _to_egg_action(self) -> bindings.Delete:
+    def _to_egg_action(self) -> bindings.Change:
         egg_call = self._call.__egg__
         if not isinstance(egg_call, bindings.Call):
             raise ValueError(f"Can only create a call with a call for the lhs, got {self._call}")  # noqa: TRY004
-        return bindings.Delete(egg_call.name, egg_call.args)
+        change: bindings._Change = bindings.Delete() if self.change == "delete" else bindings.Subsume()
+        return bindings.Change(change, egg_call.name, egg_call.args)
 
     @property
     def __egg_decls__(self) -> Declarations:
@@ -1800,16 +1803,16 @@ class Sequence(Schedule):
 
 @deprecated("Use <ruleset>.register(<rewrite>) instead of passing rulesets as arguments to rewrites.")
 @overload
-def rewrite(lhs: EXPR, ruleset: Ruleset) -> _RewriteBuilder[EXPR]: ...
+def rewrite(lhs: EXPR, ruleset: Ruleset, *, subsume: bool = False) -> _RewriteBuilder[EXPR]: ...
 
 
 @overload
-def rewrite(lhs: EXPR, ruleset: None = None) -> _RewriteBuilder[EXPR]: ...
+def rewrite(lhs: EXPR, ruleset: None = None, *, subsume: bool = False) -> _RewriteBuilder[EXPR]: ...
 
 
-def rewrite(lhs: EXPR, ruleset: Ruleset | None = None) -> _RewriteBuilder[EXPR]:
+def rewrite(lhs: EXPR, ruleset: Ruleset | None = None, *, subsume: bool = False) -> _RewriteBuilder[EXPR]:
     """Rewrite the given expression to a new expression."""
-    return _RewriteBuilder(lhs, ruleset)
+    return _RewriteBuilder(lhs, ruleset, subsume)
 
 
 @deprecated("Use <ruleset>.register(<birewrite>) instead of passing rulesets as arguments to birewrites.")
@@ -1852,7 +1855,12 @@ def expr_action(expr: Expr) -> Action:
 
 def delete(expr: Expr) -> Action:
     """Create a delete expression."""
-    return Delete(to_runtime_expr(expr))
+    return Change("delete", to_runtime_expr(expr))
+
+
+def subsume(expr: Expr) -> Action:
+    """Subsume an expression."""
+    return Change("subsume", to_runtime_expr(expr))
 
 
 def expr_fact(expr: Expr) -> Fact:
@@ -1905,10 +1913,11 @@ def vars_(names: str, bound: type[EXPR]) -> Iterable[EXPR]:
 class _RewriteBuilder(Generic[EXPR]):
     lhs: EXPR
     ruleset: Ruleset | None
+    subsume: bool
 
     def to(self, rhs: EXPR, *conditions: FactLike) -> Rewrite:
         lhs = to_runtime_expr(self.lhs)
-        rule = Rewrite(self.ruleset, lhs, convert_to_same_type(rhs, lhs), _fact_likes(conditions))
+        rule = Rewrite(self.ruleset, lhs, convert_to_same_type(rhs, lhs), _fact_likes(conditions), self.subsume)
         if self.ruleset:
             self.ruleset.append(rule)
         return rule
@@ -1924,7 +1933,7 @@ class _BirewriteBuilder(Generic[EXPR]):
 
     def to(self, rhs: EXPR, *conditions: FactLike) -> Command:
         lhs = to_runtime_expr(self.lhs)
-        rule = BiRewrite(self.ruleset, lhs, convert_to_same_type(rhs, lhs), _fact_likes(conditions))
+        rule = BiRewrite(self.ruleset, lhs, convert_to_same_type(rhs, lhs), _fact_likes(conditions), False)
         if self.ruleset:
             self.ruleset.append(rule)
         return rule

egglog/exp/array_api.py

@@ -977,7 +977,7 @@ def _astype(x: NDArray, dtype: DType, i: i64):
     ]
 
 
-@function(cost=500)
+@function
 def unique_counts(x: NDArray) -> TupleNDArray: ...
 
 
@@ -1016,7 +1016,7 @@ def _abs(f: Float):
     ]
 
 
-@function(cost=100)
+@function
 def unique_inverse(x: NDArray) -> TupleNDArray: ...
 
 
@@ -1039,7 +1039,7 @@ def zeros(
 def expand_dims(x: NDArray, axis: Int = Int(0)) -> NDArray: ...
 
 
-@function(cost=100000)
+@function
 def mean(x: NDArray, axis: OptionalIntOrTuple = OptionalIntOrTuple.none, keepdims: Boolean = FALSE) -> NDArray: ...
 
 
@@ -1048,7 +1048,7 @@ def mean(x: NDArray, axis: OptionalIntOrTuple = OptionalIntOrTuple.none, keepdim
 def sqrt(x: NDArray) -> NDArray: ...
 
 
-@function(cost=100000)
+@function
 def std(x: NDArray, axis: OptionalIntOrTuple = OptionalIntOrTuple.none) -> NDArray: ...
 
 

egglog/exp/array_api_numba.py

@@ -10,12 +10,9 @@ from egglog.exp.array_api import *
 
 array_api_numba_ruleset = ruleset()
 array_api_numba_schedule = (array_api_ruleset + array_api_numba_ruleset).saturate()
-# For these rules, we not only wanna rewrite, we also want to delete the original expression,
+# 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.
 
-# TODO: Try deleting instead if we support that in the future, and remove high cost
-# https://egraphs.zulipchat.com/#narrow/stream/375765-egglog/topic/replacing.20an.20expression.20with.20delete
-
 
 # Rewrite mean(x, <int>, <expand dims>) to use sum b/c numba cant do mean with axis
 # https://github.com/numba/numba/issues/1269
@@ -24,8 +21,8 @@ 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)).to(res)
-    yield rewrite(mean(x, axis, TRUE)).to(expand_dims(res, 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
@@ -34,7 +31,7 @@ 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)).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)
@@ -49,11 +46,11 @@ def count_values(x: NDArray, values: NDArray) -> 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)[Int(1)]).to(NDArray.vector(count_values(x, unique_values(x)))),
-        rewrite(count_values(x, NDArray.vector(TupleValue(v) + tv))).to(
+        rewrite(unique_counts(x)[Int(1)], subsume=True).to(NDArray.vector(count_values(x, unique_values(x)))),
+        rewrite(count_values(x, NDArray.vector(TupleValue(v) + tv)), subsume=True).to(
             TupleValue(sum(x == NDArray.scalar(v)).to_value()) + count_values(x, NDArray.vector(tv))
         ),
-        rewrite(count_values(x, NDArray.vector(TupleValue(v)))).to(
+        rewrite(count_values(x, NDArray.vector(TupleValue(v))), subsume=True).to(
             TupleValue(sum(x == NDArray.scalar(v)).to_value()),
         ),
     ]
@@ -64,7 +61,7 @@ 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)[Int(1)] == NDArray.scalar(Value.int(i))).to(
+        rewrite(unique_inverse(x)[Int(1)] == NDArray.scalar(Value.int(i)), subsume=True).to(
             x == NDArray.scalar(unique_values(x).index(TupleInt(i)))
         ),
     ]

{egglog-6.0.1.dist-info → egglog-6.1.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: egglog
-Version: 6.0.1
+Version: 6.1.0
 Classifier: Environment :: MacOS X
 Classifier: Environment :: Win32 (MS Windows)
 Classifier: Intended Audience :: Developers
@@ -24,6 +24,10 @@ Requires-Dist: ruff; extra == 'dev'
 Requires-Dist: mypy; extra == 'dev'
 Requires-Dist: anywidget[dev]; extra == 'dev'
 Requires-Dist: egglog[docs,test]; extra == 'dev'
+Requires-Dist: scikit-learn; extra == 'array'
+Requires-Dist: array_api_compat; extra == 'array'
+Requires-Dist: numba==0.59.0rc1; extra == 'array'
+Requires-Dist: llvmlite==0.42.0rc1; extra == 'array'
 Requires-Dist: pytest; extra == 'test'
 Requires-Dist: mypy; extra == 'test'
 Requires-Dist: syrupy; extra == 'test'
@@ -31,10 +35,6 @@ Requires-Dist: egglog[array]; extra == 'test'
 Requires-Dist: pytest-codspeed; extra == 'test'
 Requires-Dist: pytest-benchmark; extra == 'test'
 Requires-Dist: pytest-xdist; extra == 'test'
-Requires-Dist: scikit-learn; extra == 'array'
-Requires-Dist: array_api_compat; extra == 'array'
-Requires-Dist: numba==0.59.0rc1; extra == 'array'
-Requires-Dist: llvmlite==0.42.0rc1; extra == 'array'
 Requires-Dist: pydata-sphinx-theme; extra == 'docs'
 Requires-Dist: myst-nb; extra == 'docs'
 Requires-Dist: sphinx-autodoc-typehints; extra == 'docs'
@@ -48,8 +48,8 @@ Requires-Dist: line-profiler; extra == 'docs'
 Requires-Dist: sphinxcontrib-mermaid; extra == 'docs'
 Requires-Dist: ablog; extra == 'docs'
 Provides-Extra: dev
-Provides-Extra: test
 Provides-Extra: array
+Provides-Extra: test
 Provides-Extra: docs
 License-File: LICENSE
 Summary: e-graphs in Python built around the the egglog rust library

{egglog-6.0.1.dist-info → egglog-6.1.0.dist-info}/RECORD

@@ -1,11 +1,11 @@
-egglog-6.0.1.dist-info/METADATA,sha256=_0zzByEDW1HDQDSLUWZ35_k9psax38q85ZOf88XWeLM,3838
-egglog-6.0.1.dist-info/WHEEL,sha256=keLBtIOE7ZgLyd8Cijw37iqv72a2jhfxMdDu8Unr7zo,96
-egglog-6.0.1.dist-info/license_files/LICENSE,sha256=TfaboMVZ81Q6OUaKjU7z6uVjSlcGKclLYcOpgDbm9_s,1091
-egglog/bindings.pyi,sha256=bOaXFUKse5yCUrwbGY_TujRsTZfVUNYMNw8z-fgMyGk,11199
+egglog-6.1.0.dist-info/METADATA,sha256=1MhEh2rFwsmv5yqI6TZ2hoHfGwxUBSzFk1Z_slTV0HU,3838
+egglog-6.1.0.dist-info/WHEEL,sha256=keLBtIOE7ZgLyd8Cijw37iqv72a2jhfxMdDu8Unr7zo,96
+egglog-6.1.0.dist-info/license_files/LICENSE,sha256=TfaboMVZ81Q6OUaKjU7z6uVjSlcGKclLYcOpgDbm9_s,1091
+egglog/bindings.pyi,sha256=YPGerQlOpSO6kHyIYTUh1hnOopy4cyAKMPtljr3a1rg,11514
 egglog/builtins.py,sha256=lc7jWEjf5hUXzPRvo8pJrkGFIUBC0oLXV4U6WYXDrEw,11959
 egglog/config.py,sha256=mALVaxh7zmGrbuyzaVKVmYKcu1lF703QsKJw8AF7gSM,176
 egglog/declarations.py,sha256=JMivLe8Cwna3BsLgls9X_LU-ZQuFapyD-KkZpkf4Eq0,39062
-egglog/egraph.py,sha256=EtcXWeTg74QXSe4uK9dFmFIGfGmcfZWl53nj_FULMcI,69203
+egglog/egraph.py,sha256=y-vPuD0U1h8ajNrjy3V5q-r61j6yDtQbeOyDFaZy9eE,69677
 egglog/examples/bool.py,sha256=pWZTjfXR1cFy3KcihLBU5AF5rn83ImORlhUUJ1YiAXc,733
 egglog/examples/eqsat_basic.py,sha256=ORXFYYEDsEZK2IPhHtoFsd-LdjMiQi1nn7kix4Nam0s,1011
 egglog/examples/fib.py,sha256=wAn-PjazxgHDkXAU4o2xTk_GtM_iGL0biV66vWM1st4,520
@@ -16,9 +16,9 @@ egglog/examples/README.rst,sha256=QrbfmivODBvUvmY3-dHarcbC6bEvwoqAfTDhiI-aJxU,23
 egglog/examples/resolution.py,sha256=sKkbRI_v9XkQM0DriacKLINqKKDqYGFhvMCAS9tZbTA,2203
 egglog/examples/schedule_demo.py,sha256=iJtIbcLaZ7zK8UalY0z7KAKMqYjQx0MKTsNF24lKtik,652
 egglog/examples/__init__.py,sha256=KuhaJFOyz_rpUvEqZubsgLnv6rhQNE_AVFXA6bUnpdY,34
-egglog/exp/array_api.py,sha256=SlClQ4aqXL3rwWG7j-mfjbLjduA83qginlm0p7c01zM,40947
+egglog/exp/array_api.py,sha256=ozOPzHzWBT6-5TFX4Flc9Em4I_WBywZZsekmadi06wg,40901
 egglog/exp/array_api_jit.py,sha256=HIZzd0G17u-u_F4vfRdhoYvRo-ETx5HFO3RBcOfLcco,1287
-egglog/exp/array_api_numba.py,sha256=oAO9o5LfBlASDXqovdBUX0CT8GCCk8B4P9TIlBQPSps,2981
+egglog/exp/array_api_numba.py,sha256=b7pA3Mk0s0H5s0N5ZWofcWUXb4EVmnUZMEJk0EXyukw,2881
 egglog/exp/array_api_program_gen.py,sha256=crgUYXXNhQdfTq31FSIpWLIfzNsgQD8ngg3OosCtIgg,19680
 egglog/exp/program_gen.py,sha256=2qlfc-dVSNbPFQVl0eIzUzf-yw28AeaARa4PUC1xtRA,12257
 egglog/exp/__init__.py,sha256=G9zeKUcPBgIhgUg1meC86OfZVFETYIugyHWseTcCe_0,52
@@ -30,5 +30,5 @@ egglog/type_constraint_solver.py,sha256=u-b2SZxptCxjYNcht_TQSAebIVEcf3P_-2-TJk8u
 egglog/widget.css,sha256=WJS2M1wQdujhSTCakMa_ZXuoTPre1Uy1lPcvBE1LZQU,102
 egglog/widget.js,sha256=UNOER3sYZ-bS7Qhw9S6qtpR81FuHa5DzXQaWWtQq334,2021
 egglog/__init__.py,sha256=AGM-7H_Vx1Xb0Z-tc0C9Y8JbQxpa6hdYfhSU_dnT3Mw,244
-egglog/bindings.cp310-win_amd64.pyd,sha256=rqr6-ruHokEX0FshXaGf6dN4-GLH3l6KFcSbVP0Kp70,4878336
-egglog-6.0.1.dist-info/RECORD,,
+egglog/bindings.cp310-win_amd64.pyd,sha256=M6YHB0Ef-G-xSpFHmIFtPa-sL8yQ2t_C_DfDoPjCJ0s,4991488
+egglog-6.1.0.dist-info/RECORD,,