diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index dca3c646..0bf544b6 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,12 +1,12 @@ exclude: ^python/tests/__snapshots__/ repos: - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.14.3 + rev: v0.16.5 hooks: - id: ruff-check args: [--fix] - id: ruff-format - repo: https://github.com/astral-sh/uv-pre-commit - rev: 0.9.7 + rev: 0.12.7 hooks: - id: uv-lock diff --git a/docs/changelog.md b/docs/changelog.md index abf97873..ee455ebf 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -210,15 +210,18 @@ a linalg function (in [an example inspired by Siu](https://gist.github.com/sklam ```python from egglog.exp.array_api import * + @function(ruleset=array_api_ruleset, subsume=True) def linalg_norm(X: NDArrayLike, axis: TupleIntLike) -> NDArray: X = cast(NDArray, X) return NDArray( X.shape.deselect(axis), X.dtype, - lambda k: ndindex(X.shape.select(axis)) - .foldl_value(lambda carry, i: carry + ((x := X.index(i + k)).conj() * x).real(), init=0.0) - .sqrt(), + lambda k: ( + ndindex(X.shape.select(axis)) + .foldl_value(lambda carry, i: carry + ((x := X.index(i + k)).conj() * x).real(), init=0.0) + .sqrt() + ), ) ``` @@ -372,8 +375,8 @@ rule or expression. For example: class A(Expr): def __init__(self, b: B) -> None: ... -class B(Expr): - ... + +class B(Expr): ... ``` ### Top level commands diff --git a/docs/explanation/2026_02_containers.md b/docs/explanation/2026_02_containers.md index e504eea1..6a211a5e 100644 --- a/docs/explanation/2026_02_containers.md +++ b/docs/explanation/2026_02_containers.md @@ -515,6 +515,7 @@ def get_sole_polynomial(xs: MultiSet[Value]) -> MultiSet[MultiSet[Value]]: get_sole_polynomial(MultiSet(polynomial(xss))) => xss """ + @ruleset def to_polynomial_ruleset( n1: Value, @@ -560,29 +561,29 @@ nested. For example, we might end up with a term like `polynomial(MultiSet(Multi with just `polynomial(xs)`. We define two additional rules to cover cases like this: ```python - yield rule( - eq(n1).to(polynomial(mss)), - # For each monomial, if any of its terms is a polynomial with a single monomial, flatten - # that into the monomial, otherwise keep it as is - mss1 == mss.map(partial(multiset_flat_map, get_monomial)), - mss != mss1, # skip if this is a no-op - name="unwrap monomial", - ).then( - union(n1).with_(polynomial(mss1)), - delete(polynomial(mss)), - set_(get_sole_polynomial(MultiSet(polynomial(mss1)))).to(mss1), - ) - yield rule( - eq(n1).to(polynomial(mss)), - # If any of the monomials just has a single item which is a polynomial, then flatten that into the outer polynomial - mss1 == multiset_flat_map(UnstableFn(get_sole_polynomial), mss), - mss != mss1, - name="unwrap polynomial", - ).then( - union(n1).with_(polynomial(mss1)), - delete(polynomial(mss)), - set_(get_sole_polynomial(MultiSet(polynomial(mss1)))).to(mss1), - ) +yield rule( + eq(n1).to(polynomial(mss)), + # For each monomial, if any of its terms is a polynomial with a single monomial, flatten + # that into the monomial, otherwise keep it as is + mss1 == mss.map(partial(multiset_flat_map, get_monomial)), + mss != mss1, # skip if this is a no-op + name="unwrap monomial", +).then( + union(n1).with_(polynomial(mss1)), + delete(polynomial(mss)), + set_(get_sole_polynomial(MultiSet(polynomial(mss1)))).to(mss1), +) +yield rule( + eq(n1).to(polynomial(mss)), + # If any of the monomials just has a single item which is a polynomial, then flatten that into the outer polynomial + mss1 == multiset_flat_map(UnstableFn(get_sole_polynomial), mss), + mss != mss1, + name="unwrap polynomial", +).then( + union(n1).with_(polynomial(mss1)), + delete(polynomial(mss)), + set_(get_sole_polynomial(MultiSet(polynomial(mss1)))).to(mss1), +) ``` We have avoided the need to match inside of containers by instead using higher order functions to apply blockwise @@ -643,12 +644,12 @@ def factor_ruleset( eq(n).to(polynomial(mss)), # Find factor that shows up in most monomials, at least two of them counts == MultiSet.sum_multisets(mss.map(MultiSet.reset_counts)), - eq(picked_term).to(counts.pick_max()), # on ties pick an arbitrary one + eq(picked_term).to(counts.pick_max()), # on ties pick an arbitrary one # Only factor out if it appears in more than one monomial counts.count(picked_term) > 1, # The factor we choose is the largest intersection between all the monomials that have that that factored term picked == mss.filter(partial(multiset_contains_swapped, picked_term)), - factor == multiset_fold(MultiSet.__and__, picked.pick(), picked), # intersection + factor == multiset_fold(MultiSet.__and__, picked.pick(), picked), # intersection divided == picked.map(partial(multiset_subtract_swapped, factor)), # remainder is those monomials that do not contain the factor remainder == mss.filter(partial(multiset_not_contains_swapped, picked_term)), diff --git a/docs/explanation/optional_values.md b/docs/explanation/optional_values.md index 21e88578..d9a0c30d 100644 --- a/docs/explanation/optional_values.md +++ b/docs/explanation/optional_values.md @@ -28,14 +28,9 @@ Let's see what happens if we call check on something that is not true: ```python from egglog.bindings import * -egraph = EGraph() -egraph.check_fact(Fact( - Call( - "<", - [Lit(Int(2)), Lit(Int(1))] - ) -)) +egraph = EGraph() +egraph.check_fact(Fact(Call("<", [Lit(Int(2)), Lit(Int(1))]))) ``` We get: @@ -48,14 +43,9 @@ And what if we call check with a non-unit value? ```python from egglog.bindings import * -egraph = EGraph() -egraph.check_fact(Fact( - Call( - "+", - [Lit(Int(2)), Lit(Int(1))] - ) -)) +egraph = EGraph() +egraph.check_fact(Fact(Call("+", [Lit(Int(2)), Lit(Int(1))]))) ``` Yep, it fails on us: diff --git a/docs/reference/python-integration.md b/docs/reference/python-integration.md index 23f52d8f..1cf84b0b 100644 --- a/docs/reference/python-integration.md +++ b/docs/reference/python-integration.md @@ -85,6 +85,7 @@ class MyExpr(Expr): return value raise ExprValueError(self, "MyExpr") + match MyExpr("hello"): case MyExpr(value): print(f"Matched MyExpr with value: {value}") @@ -602,7 +603,7 @@ egraph This is equivalent to adding the rewrite rules to the e-graph directly, like this, but just more succinct: ```python -x = var("x", Math) +x = var("x", Math) egraph.register(rewrite(pi).to(math_float(3.14))) egraph.register(rewrite(square(x)).to(x * x)) ``` diff --git a/python/egglog/__init__.py b/python/egglog/__init__.py index 27b2bc61..5f937c99 100644 --- a/python/egglog/__init__.py +++ b/python/egglog/__init__.py @@ -4,7 +4,7 @@ from . import config, ipython_magic # noqa: F401 from .bindings import EggSmolError, StageInfo, TimeOnly, WithPlan # noqa: F401 -from .builtins import * # noqa: UP029 +from .builtins import * from .conversion import * from .deconstruct import * from .egraph import * diff --git a/python/egglog/declarations.py b/python/egglog/declarations.py index 9a5d08e2..5e53cf42 100644 --- a/python/egglog/declarations.py +++ b/python/egglog/declarations.py @@ -49,7 +49,6 @@ "ConstantRef", "ConstructorDecl", "Declarations", - "Declarations", "DeclarationsLike", "DefaultRewriteDecl", "DelayedDeclarations", @@ -128,7 +127,7 @@ class HasDeclarations(Protocol): def __egg_decls__(self) -> Declarations: ... -DeclarationsLike: TypeAlias = Union[HasDeclarations, None, "Declarations"] +DeclarationsLike: TypeAlias = Union[HasDeclarations, "Declarations", None] def upcast_declarations(declarations_like: Iterable[DeclarationsLike]) -> list[Declarations]: @@ -659,7 +658,7 @@ def signature(self) -> FunctionSignature: ) @property - def egg_name(self) -> None | str: + def egg_name(self) -> str | None: return None diff --git a/python/egglog/egraph.py b/python/egglog/egraph.py index 6fe782d5..e9334371 100644 --- a/python/egglog/egraph.py +++ b/python/egglog/egraph.py @@ -60,13 +60,11 @@ "BaseExpr", "BuiltinExpr", "Command", - "Command", "CostModel", "EGraph", "Expr", "ExprCallable", "Fact", - "Fact", "GraphvizKwargs", "GreedyDagCost", "RewriteOrRule", @@ -1538,7 +1536,7 @@ def __call__(self, *args, **kwargs) -> Never: def ruleset( rule_or_generator: RewriteOrRule | RewriteOrRuleGenerator | None = None, *rules: RewriteOrRule, - name: None | str = None, + name: str | None = None, ) -> Ruleset: """ Creates a ruleset with the following rules. @@ -2086,7 +2084,7 @@ def run(ruleset: Ruleset | None = None, *until: FactLike, scheduler: BackOff | N ) -def back_off(match_limit: None | int = None, ban_length: None | int = None) -> BackOff: +def back_off(match_limit: int | None = None, ban_length: int | None = None) -> BackOff: """ Create a backoff scheduler configuration. diff --git a/python/egglog/exp/any_expr.py b/python/egglog/exp/any_expr.py index 4de0d3ff..07eb8097 100644 --- a/python/egglog/exp/any_expr.py +++ b/python/egglog/exp/any_expr.py @@ -753,8 +753,8 @@ def any_eval(self: A) -> Any: return res -_CURRENT_EGRAPH: None | EGraph = None -_LAST_ASSERT: None | A = None +_CURRENT_EGRAPH: EGraph | None = None +_LAST_ASSERT: A | None = None @contextlib.contextmanager diff --git a/python/egglog/exp/array_api.py b/python/egglog/exp/array_api.py index 42ea19a0..ba4d6290 100644 --- a/python/egglog/exp/array_api.py +++ b/python/egglog/exp/array_api.py @@ -1459,7 +1459,7 @@ def slice(cls, slice: Slice) -> MultiAxisIndexKeyItem: ... converter(Int, MultiAxisIndexKeyItem, lambda i: MultiAxisIndexKeyItem.int(i)) converter(Slice, MultiAxisIndexKeyItem, lambda s: MultiAxisIndexKeyItem.slice(s)) -MultiAxisIndexKeyItemLike: TypeAlias = MultiAxisIndexKeyItem | EllipsisType | None | IntLike | SliceLike +MultiAxisIndexKeyItemLike: TypeAlias = MultiAxisIndexKeyItem | EllipsisType | IntLike | SliceLike | None class MultiAxisIndexKey(Expr, ruleset=array_api_ruleset): @@ -2123,7 +2123,7 @@ def int(cls, value: Int) -> OptionalIntOrTuple: ... def tuple(cls, value: TupleIntLike) -> OptionalIntOrTuple: ... -OptionalIntOrTupleLike: TypeAlias = OptionalIntOrTuple | None | IntLike | TupleIntLike +OptionalIntOrTupleLike: TypeAlias = OptionalIntOrTuple | IntLike | TupleIntLike | None converter(type(None), OptionalIntOrTuple, lambda _: OptionalIntOrTuple.none) converter(Int, OptionalIntOrTuple, lambda v: OptionalIntOrTuple.int(v)) @@ -2379,7 +2379,8 @@ def vecdot(x1: NDArrayLike, x2: NDArrayLike) -> NDArray: x1.shape.drop_last(), x1.dtype, lambda idx: ( - TupleInt.range(x1.shape.last()) + TupleInt + .range(x1.shape.last()) .map_value(lambda i: x1.index(idx.append(i)) * x2.index((i,))) .foldl_value(Value.__add__, Value.from_float(0)) ), @@ -2740,7 +2741,8 @@ def unravel_index(flat_index: IntLike, shape: TupleIntLike) -> TupleInt: shape = cast("TupleInt", shape) return ( - shape.reverse() + shape + .reverse() .foldl_tuple_int( # Store remainder as last item in accumulator lambda acc, dim: acc.drop_last().append((r := acc.last()) % dim).append(r // dim), @@ -2754,7 +2756,7 @@ def unravel_index(flat_index: IntLike, shape: TupleIntLike) -> TupleInt: array_api_combined_ruleset = array_api_ruleset array_api_schedule = (array_api_combined_ruleset + run()).saturate() -_CURRENT_EGRAPH: None | EGraph = None +_CURRENT_EGRAPH: EGraph | None = None @contextlib.contextmanager diff --git a/python/egglog/runtime.py b/python/egglog/runtime.py index 2ec5a453..bebe9c55 100644 --- a/python/egglog/runtime.py +++ b/python/egglog/runtime.py @@ -649,7 +649,7 @@ def to_py_signature(sig: FunctionSignature, decls: Declarations, optional_args: return Signature(parameters) -ON_CREATE_EXPR: None | Callable[[Callable[[], TypedExprDecl]], None] = None +ON_CREATE_EXPR: Callable[[Callable[[], TypedExprDecl]], None] | None = None @dataclass diff --git a/python/tests/test_array_api.py b/python/tests/test_array_api.py index d9eea325..e401ae39 100644 --- a/python/tests/test_array_api.py +++ b/python/tests/test_array_api.py @@ -263,11 +263,14 @@ def linalg_norm(X: NDArray, axis: TupleIntLike) -> NDArray: return NDArray.fn( outshape, X.dtype, - lambda k: LoopNestAPI.from_tuple(reduce_axis) - .unwrap() - .indices() - .foldl_value(lambda carry, i: carry + ((x := X.index(i + k)).conj() * x).real(), init=0.0) - .sqrt(), + lambda k: ( + LoopNestAPI + .from_tuple(reduce_axis) + .unwrap() + .indices() + .foldl_value(lambda carry, i: carry + ((x := X.index(i + k)).conj() * x).real(), init=0.0) + .sqrt() + ), ) @@ -277,9 +280,11 @@ def linalg_norm_v2(X: NDArrayLike, axis: TupleIntLike) -> NDArray: return NDArray.fn( X.shape.deselect(axis), X.dtype, - lambda k: ndindex(X.shape.select(axis)) - .foldl_value(lambda carry, i: carry + ((x := X.index(i + k)).conj() * x).real(), init=0.0) - .sqrt(), + lambda k: ( + ndindex(X.shape.select(axis)) + .foldl_value(lambda carry, i: carry + ((x := X.index(i + k)).conj() * x).real(), init=0.0) + .sqrt() + ), )