Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions .pre-commit-config.yaml
Original file line number Diff line number Diff line change
@@ -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
13 changes: 8 additions & 5 deletions docs/changelog.md
Original file line number Diff line number Diff line change
Expand Up @@ -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()
),
)
```

Expand Down Expand Up @@ -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
Expand Down
51 changes: 26 additions & 25 deletions docs/explanation/2026_02_containers.md
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)),
Expand Down
18 changes: 4 additions & 14 deletions docs/explanation/optional_values.md
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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:
Expand Down
3 changes: 2 additions & 1 deletion docs/reference/python-integration.md
Original file line number Diff line number Diff line change
Expand Up @@ -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}")
Expand Down Expand Up @@ -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))
```
Expand Down
2 changes: 1 addition & 1 deletion python/egglog/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 *
Expand Down
5 changes: 2 additions & 3 deletions python/egglog/declarations.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,6 @@
"ConstantRef",
"ConstructorDecl",
"Declarations",
"Declarations",
"DeclarationsLike",
"DefaultRewriteDecl",
"DelayedDeclarations",
Expand Down Expand Up @@ -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]:
Expand Down Expand Up @@ -659,7 +658,7 @@ def signature(self) -> FunctionSignature:
)

@property
def egg_name(self) -> None | str:
def egg_name(self) -> str | None:
return None


Expand Down
6 changes: 2 additions & 4 deletions python/egglog/egraph.py
Original file line number Diff line number Diff line change
Expand Up @@ -60,13 +60,11 @@
"BaseExpr",
"BuiltinExpr",
"Command",
"Command",
"CostModel",
"EGraph",
"Expr",
"ExprCallable",
"Fact",
"Fact",
"GraphvizKwargs",
"GreedyDagCost",
"RewriteOrRule",
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.

Expand Down
4 changes: 2 additions & 2 deletions python/egglog/exp/any_expr.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
12 changes: 7 additions & 5 deletions python/egglog/exp/array_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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))
Expand Down Expand Up @@ -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))
),
Expand Down Expand Up @@ -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),
Expand All @@ -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
Expand Down
2 changes: 1 addition & 1 deletion python/egglog/runtime.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
21 changes: 13 additions & 8 deletions python/tests/test_array_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
),
)


Expand All @@ -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()
),
)


Expand Down
Loading