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

egglog/egraph_state.py

@@ -0,0 +1,417 @@
+"""
+Implement conversion to/from egglog.
+"""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, overload
+from weakref import WeakKeyDictionary
+
+from typing_extensions import assert_never
+
+from . import bindings
+from .declarations import *
+from .pretty import *
+from .type_constraint_solver import TypeConstraintError, TypeConstraintSolver
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+
+__all__ = ["EGraphState", "GLOBAL_PY_OBJECT_SORT"]
+
+# Create a global sort for python objects, so we can store them without an e-graph instance
+# Needed when serializing commands to egg commands when creating modules
+GLOBAL_PY_OBJECT_SORT = bindings.PyObjectSort()
+
+
+@dataclass
+class EGraphState:
+    """
+    State of the EGraph declerations and rulesets, so when we pop/push the stack we know whats defined.
+
+    Used for converting to/from egg and for pretty printing.
+    """
+
+    egraph: bindings.EGraph
+    # The decleratons we have added.
+    __egg_decls__: Declarations = field(default_factory=Declarations)
+    # Mapping of added rulesets to the added rules
+    rulesets: dict[str, set[RewriteOrRuleDecl]] = field(default_factory=dict)
+
+    # Bidirectional mapping between egg function names and python callable references.
+    # Note that there are possibly mutliple callable references for a single egg function name, like `+`
+    # for both int and rational classes.
+    egg_fn_to_callable_refs: dict[str, set[CallableRef]] = field(
+        default_factory=lambda: defaultdict(set, {"!=": {FunctionRef("!=")}})
+    )
+    callable_ref_to_egg_fn: dict[CallableRef, str] = field(default_factory=lambda: {FunctionRef("!="): "!="})
+
+    # Bidirectional mapping between egg sort names and python type references.
+    type_ref_to_egg_sort: dict[JustTypeRef, str] = field(default_factory=dict)
+
+    # Cache of egg expressions for converting to egg
+    expr_to_egg_cache: WeakKeyDictionary[ExprDecl, bindings._Expr] = field(default_factory=WeakKeyDictionary)
+
+    def copy(self) -> EGraphState:
+        """
+        Returns a copy of the state. Th egraph reference is kept the same. Used for pushing/popping.
+        """
+        return EGraphState(
+            egraph=self.egraph,
+            __egg_decls__=self.__egg_decls__.copy(),
+            rulesets={k: v.copy() for k, v in self.rulesets.items()},
+            egg_fn_to_callable_refs=defaultdict(set, {k: v.copy() for k, v in self.egg_fn_to_callable_refs.items()}),
+            callable_ref_to_egg_fn=self.callable_ref_to_egg_fn.copy(),
+            type_ref_to_egg_sort=self.type_ref_to_egg_sort.copy(),
+            expr_to_egg_cache=self.expr_to_egg_cache.copy(),
+        )
+
+    def schedule_to_egg(self, schedule: ScheduleDecl) -> bindings._Schedule:
+        match schedule:
+            case SaturateDecl(schedule):
+                return bindings.Saturate(self.schedule_to_egg(schedule))
+            case RepeatDecl(schedule, times):
+                return bindings.Repeat(times, self.schedule_to_egg(schedule))
+            case SequenceDecl(schedules):
+                return bindings.Sequence([self.schedule_to_egg(s) for s in schedules])
+            case RunDecl(ruleset_name, until):
+                self.ruleset_to_egg(ruleset_name)
+                config = bindings.RunConfig(ruleset_name, None if not until else list(map(self.fact_to_egg, until)))
+                return bindings.Run(config)
+            case _:
+                assert_never(schedule)
+
+    def ruleset_to_egg(self, name: str) -> None:
+        """
+        Registers a ruleset if it's not already registered.
+        """
+        if name not in self.rulesets:
+            if name:
+                self.egraph.run_program(bindings.AddRuleset(name))
+            rules = self.rulesets[name] = set()
+        else:
+            rules = self.rulesets[name]
+        for rule in self.__egg_decls__._rulesets[name].rules:
+            if rule in rules:
+                continue
+            self.egraph.run_program(self.command_to_egg(rule, name))
+            rules.add(rule)
+
+    def command_to_egg(self, cmd: CommandDecl, ruleset: str) -> bindings._Command:
+        match cmd:
+            case ActionCommandDecl(action):
+                return bindings.ActionCommand(self.action_to_egg(action))
+            case RewriteDecl(tp, lhs, rhs, conditions) | BiRewriteDecl(tp, lhs, rhs, conditions):
+                self.type_ref_to_egg(tp)
+                rewrite = bindings.Rewrite(
+                    self.expr_to_egg(lhs),
+                    self.expr_to_egg(rhs),
+                    [self.fact_to_egg(c) for c in conditions],
+                )
+                return (
+                    bindings.RewriteCommand(ruleset, rewrite, cmd.subsume)
+                    if isinstance(cmd, RewriteDecl)
+                    else bindings.BiRewriteCommand(ruleset, rewrite)
+                )
+            case RuleDecl(head, body, name):
+                rule = bindings.Rule(
+                    [self.action_to_egg(a) for a in head],
+                    [self.fact_to_egg(f) for f in body],
+                )
+                return bindings.RuleCommand(name or "", ruleset, rule)
+            case _:
+                assert_never(cmd)
+
+    def action_to_egg(self, action: ActionDecl) -> bindings._Action:
+        match action:
+            case LetDecl(name, typed_expr):
+                return bindings.Let(name, self.typed_expr_to_egg(typed_expr))
+            case SetDecl(tp, call, rhs):
+                self.type_ref_to_egg(tp)
+                call_ = self.expr_to_egg(call)
+                return bindings.Set(call_.name, call_.args, self.expr_to_egg(rhs))
+            case ExprActionDecl(typed_expr):
+                return bindings.Expr_(self.typed_expr_to_egg(typed_expr))
+            case ChangeDecl(tp, call, change):
+                self.type_ref_to_egg(tp)
+                call_ = self.expr_to_egg(call)
+                egg_change: bindings._Change
+                match change:
+                    case "delete":
+                        egg_change = bindings.Delete()
+                    case "subsume":
+                        egg_change = bindings.Subsume()
+                    case _:
+                        assert_never(change)
+                return bindings.Change(egg_change, call_.name, call_.args)
+            case UnionDecl(tp, lhs, rhs):
+                self.type_ref_to_egg(tp)
+                return bindings.Union(self.expr_to_egg(lhs), self.expr_to_egg(rhs))
+            case PanicDecl(name):
+                return bindings.Panic(name)
+            case _:
+                assert_never(action)
+
+    def fact_to_egg(self, fact: FactDecl) -> bindings._Fact:
+        match fact:
+            case EqDecl(tp, exprs):
+                self.type_ref_to_egg(tp)
+                return bindings.Eq([self.expr_to_egg(e) for e in exprs])
+            case ExprFactDecl(typed_expr):
+                return bindings.Fact(self.typed_expr_to_egg(typed_expr))
+            case _:
+                assert_never(fact)
+
+    def callable_ref_to_egg(self, ref: CallableRef) -> str:
+        """
+        Returns the egg function name for a callable reference, registering it if it is not already registered.
+        """
+        if ref in self.callable_ref_to_egg_fn:
+            return self.callable_ref_to_egg_fn[ref]
+        decl = self.__egg_decls__.get_callable_decl(ref)
+        self.callable_ref_to_egg_fn[ref] = egg_name = decl.egg_name or _generate_callable_egg_name(ref)
+        self.egg_fn_to_callable_refs[egg_name].add(ref)
+        match decl:
+            case RelationDecl(arg_types, _, _):
+                self.egraph.run_program(bindings.Relation(egg_name, [self.type_ref_to_egg(a) for a in arg_types]))
+            case ConstantDecl(tp, _):
+                # Use function decleration instead of constant b/c constants cannot be extracted
+                # https://github.com/egraphs-good/egglog/issues/334
+                self.egraph.run_program(
+                    bindings.Function(bindings.FunctionDecl(egg_name, bindings.Schema([], self.type_ref_to_egg(tp))))
+                )
+            case FunctionDecl():
+                if not decl.builtin:
+                    egg_fn_decl = bindings.FunctionDecl(
+                        egg_name,
+                        bindings.Schema(
+                            [self.type_ref_to_egg(a.to_just()) for a in decl.arg_types],
+                            self.type_ref_to_egg(decl.semantic_return_type.to_just()),
+                        ),
+                        self.expr_to_egg(decl.default) if decl.default else None,
+                        self.expr_to_egg(decl.merge) if decl.merge else None,
+                        [self.action_to_egg(a) for a in decl.on_merge],
+                        decl.cost,
+                        decl.unextractable,
+                    )
+                    self.egraph.run_program(bindings.Function(egg_fn_decl))
+            case _:
+                assert_never(decl)
+        return egg_name
+
+    def type_ref_to_egg(self, ref: JustTypeRef) -> str:
+        """
+        Returns the egg sort name for a type reference, registering it if it is not already registered.
+        """
+        try:
+            return self.type_ref_to_egg_sort[ref]
+        except KeyError:
+            pass
+        decl = self.__egg_decls__._classes[ref.name]
+        self.type_ref_to_egg_sort[ref] = egg_name = decl.egg_name or _generate_type_egg_name(ref)
+        if not decl.builtin or ref.args:
+            self.egraph.run_program(
+                bindings.Sort(
+                    egg_name,
+                    (
+                        (
+                            self.type_ref_to_egg(JustTypeRef(ref.name)),
+                            [bindings.Var(self.type_ref_to_egg(a)) for a in ref.args],
+                        )
+                        if ref.args
+                        else None
+                    ),
+                )
+            )
+        # For builtin classes, let's also make sure we have the mapping of all egg fn names for class methods, because
+        # these can be created even without adding them to the e-graph, like `vec-empty` which can be extracted
+        # even if you never use that function.
+        if decl.builtin:
+            for method in decl.class_methods:
+                self.callable_ref_to_egg(ClassMethodRef(ref.name, method))
+
+        return egg_name
+
+    def op_mapping(self) -> dict[str, str]:
+        """
+        Create a mapping of egglog function name to Python function name, for use in the serialized format
+        for better visualization.
+        """
+        return {
+            k: pretty_callable_ref(self.__egg_decls__, next(iter(v)))
+            for k, v in self.egg_fn_to_callable_refs.items()
+            if len(v) == 1
+        }
+
+    def typed_expr_to_egg(self, typed_expr_decl: TypedExprDecl) -> bindings._Expr:
+        self.type_ref_to_egg(typed_expr_decl.tp)
+        return self.expr_to_egg(typed_expr_decl.expr)
+
+    @overload
+    def expr_to_egg(self, expr_decl: CallDecl) -> bindings.Call: ...
+
+    @overload
+    def expr_to_egg(self, expr_decl: ExprDecl) -> bindings._Expr: ...
+
+    def expr_to_egg(self, expr_decl: ExprDecl) -> bindings._Expr:
+        """
+        Convert an ExprDecl to an egg expression.
+
+        Cached using weakrefs to avoid memory leaks.
+        """
+        try:
+            return self.expr_to_egg_cache[expr_decl]
+        except KeyError:
+            pass
+
+        res: bindings._Expr
+        match expr_decl:
+            case VarDecl(name):
+                res = bindings.Var(name)
+            case LitDecl(value):
+                l: bindings._Literal
+                match value:
+                    case None:
+                        l = bindings.Unit()
+                    case bool(i):
+                        l = bindings.Bool(i)
+                    case int(i):
+                        l = bindings.Int(i)
+                    case float(f):
+                        l = bindings.F64(f)
+                    case str(s):
+                        l = bindings.String(s)
+                    case _:
+                        assert_never(value)
+                res = bindings.Lit(l)
+            case CallDecl(ref, args, _):
+                egg_fn = self.callable_ref_to_egg(ref)
+                egg_args = [self.typed_expr_to_egg(a) for a in args]
+                res = bindings.Call(egg_fn, egg_args)
+            case PyObjectDecl(value):
+                res = GLOBAL_PY_OBJECT_SORT.store(value)
+            case _:
+                assert_never(expr_decl.expr)
+
+        self.expr_to_egg_cache[expr_decl] = res
+        return res
+
+    def exprs_from_egg(
+        self, termdag: bindings.TermDag, terms: list[bindings._Term], tp: JustTypeRef
+    ) -> Iterable[TypedExprDecl]:
+        """
+        Create a function that can convert from an egg term to a typed expr.
+        """
+        state = FromEggState(self, termdag)
+        return [state.from_expr(tp, term) for term in terms]
+
+    def _get_possible_types(self, cls_name: str) -> frozenset[JustTypeRef]:
+        """
+        Given a class name, returns all possible registered types that it can be.
+        """
+        return frozenset(tp for tp in self.type_ref_to_egg_sort if tp.name == cls_name)
+
+
+def _generate_type_egg_name(ref: JustTypeRef) -> str:
+    """
+    Generates an egg sort name for this type reference by linearizing the type.
+    """
+    name = ref.name
+    if not ref.args:
+        return name
+    return f"{name}_{'_'.join(map(_generate_type_egg_name, ref.args))}"
+
+
+def _generate_callable_egg_name(ref: CallableRef) -> str:
+    """
+    Generates a valid egg function name for a callable reference.
+    """
+    match ref:
+        case FunctionRef(name) | ConstantRef(name):
+            return name
+        case (
+            MethodRef(cls_name, name)
+            | ClassMethodRef(cls_name, name)
+            | ClassVariableRef(cls_name, name)
+            | PropertyRef(cls_name, name)
+        ):
+            return f"{cls_name}_{name}"
+        case _:
+            assert_never(ref)
+
+
+@dataclass
+class FromEggState:
+    """
+    Dataclass containing state used when converting from an egg term to a typed expr.
+    """
+
+    state: EGraphState
+    termdag: bindings.TermDag
+    # Cache of termdag ID to TypedExprDecl
+    cache: dict[int, TypedExprDecl] = field(default_factory=dict)
+
+    @property
+    def decls(self) -> Declarations:
+        return self.state.__egg_decls__
+
+    def from_expr(self, tp: JustTypeRef, term: bindings._Term) -> TypedExprDecl:
+        """
+        Convert an egg term to a typed expr.
+        """
+        expr_decl: ExprDecl
+        if isinstance(term, bindings.TermVar):
+            expr_decl = VarDecl(term.name)
+        elif isinstance(term, bindings.TermLit):
+            value = term.value
+            expr_decl = LitDecl(None if isinstance(value, bindings.Unit) else value.value)
+        elif isinstance(term, bindings.TermApp):
+            if term.name == "py-object":
+                call = bindings.termdag_term_to_expr(self.termdag, term)
+                expr_decl = PyObjectDecl(self.state.egraph.eval_py_object(call))
+            else:
+                expr_decl = self.from_call(tp, term)
+        else:
+            assert_never(term)
+        return TypedExprDecl(tp, expr_decl)
+
+    def from_call(self, tp: JustTypeRef, term: bindings.TermApp) -> CallDecl:
+        """
+        Convert a call to a CallDecl.
+
+        There could be Python call refs which match the call, so we need to find the correct one.
+        """
+        # Find the first callable ref that matches the call
+        for callable_ref in self.state.egg_fn_to_callable_refs[term.name]:
+            # If this is a classmethod, we might need the type params that were bound for this type
+            # This could be multiple types if the classmethod is ambiguous, like map create.
+            possible_types: Iterable[JustTypeRef | None]
+            fn_decl = self.decls.get_callable_decl(callable_ref).to_function_decl()
+            if isinstance(callable_ref, ClassMethodRef):
+                possible_types = self.state._get_possible_types(callable_ref.class_name)
+                cls_name = callable_ref.class_name
+            else:
+                possible_types = [None]
+                cls_name = None
+            for possible_type in possible_types:
+                tcs = TypeConstraintSolver(self.decls)
+                if possible_type and possible_type.args:
+                    tcs.bind_class(possible_type)
+
+                try:
+                    arg_types, bound_tp_params = tcs.infer_arg_types(
+                        fn_decl.arg_types, fn_decl.semantic_return_type, fn_decl.var_arg_type, tp, cls_name
+                    )
+                except TypeConstraintError:
+                    continue
+                args: list[TypedExprDecl] = []
+                for a, tp in zip(term.args, arg_types, strict=False):
+                    try:
+                        res = self.cache[a]
+                    except KeyError:
+                        res = self.cache[a] = self.from_expr(tp, self.termdag.nodes[a])
+                    args.append(res)
+                return CallDecl(callable_ref, tuple(args), bound_tp_params)
+        raise ValueError(f"Could not find callable ref for call {term}")

egglog/exp/array_api.py

@@ -24,7 +24,7 @@ if TYPE_CHECKING:
 # Pretend that exprs are numbers b/c sklearn does isinstance checks
 numbers.Integral.register(RuntimeExpr)
 
-array_api_ruleset = ruleset()
+array_api_ruleset = ruleset(name="array_api_ruleset")
 array_api_schedule = array_api_ruleset.saturate()
 
 
@@ -36,10 +36,14 @@ class Boolean(Expr):
     @property
     def bool(self) -> Bool: ...
 
-    def __or__(self, other: Boolean) -> Boolean: ...
+    def __or__(self, other: BooleanLike) -> Boolean: ...
 
-    def __and__(self, other: Boolean) -> Boolean: ...
+    def __and__(self, other: BooleanLike) -> Boolean: ...
 
+    def if_int(self, true_value: Int, false_value: Int) -> Int: ...
+
+
+BooleanLike = Boolean | bool
 
 TRUE = constant("TRUE", Boolean)
 FALSE = constant("FALSE", Boolean)
@@ -47,7 +51,7 @@ converter(bool, Boolean, lambda x: TRUE if x else FALSE)
 
 
 @array_api_ruleset.register
-def _bool(x: Boolean):
+def _bool(x: Boolean, i: Int, j: Int):
     return [
         rule(eq(x).to(TRUE)).then(set_(x.bool).to(Bool(True))),
         rule(eq(x).to(FALSE)).then(set_(x.bool).to(Bool(False))),
@@ -55,82 +59,8 @@ def _bool(x: Boolean):
         rewrite(FALSE | x).to(x),
         rewrite(TRUE & x).to(x),
         rewrite(FALSE & x).to(FALSE),
-    ]
-
-
-class DType(Expr):
-    float64: ClassVar[DType]
-    float32: ClassVar[DType]
-    int64: ClassVar[DType]
-    int32: ClassVar[DType]
-    object: ClassVar[DType]
-    bool: ClassVar[DType]
-
-    def __eq__(self, other: DType) -> Boolean:  # type: ignore[override]
-        ...
-
-
-float64 = DType.float64
-float32 = DType.float32
-int32 = DType.int32
-int64 = DType.int64
-
-_DTYPES = [float64, float32, int32, int64, DType.object]
-
-converter(type, DType, lambda x: convert(np.dtype(x), DType))
-converter(type(np.dtype), DType, lambda x: getattr(DType, x.name))  # type: ignore[call-overload]
-array_api_ruleset.register(
-    *(rewrite(l == r).to(TRUE if l is r else FALSE) for l, r in itertools.product(_DTYPES, repeat=2))
-)
-
-
-class IsDtypeKind(Expr):
-    NULL: ClassVar[IsDtypeKind]
-
-    @classmethod
-    def string(cls, s: StringLike) -> IsDtypeKind: ...
-
-    @classmethod
-    def dtype(cls, d: DType) -> IsDtypeKind: ...
-
-    @method(cost=10)
-    def __or__(self, other: IsDtypeKind) -> IsDtypeKind: ...
-
-
-# TODO: Make kind more generic to support tuples.
-@function
-def isdtype(dtype: DType, kind: IsDtypeKind) -> Boolean: ...
-
-
-converter(DType, IsDtypeKind, lambda x: IsDtypeKind.dtype(x))
-converter(str, IsDtypeKind, lambda x: IsDtypeKind.string(x))
-converter(
-    tuple, IsDtypeKind, lambda x: convert(x[0], IsDtypeKind) | convert(x[1:], IsDtypeKind) if x else IsDtypeKind.NULL
-)
-
-
-@array_api_ruleset.register
-def _isdtype(d: DType, k1: IsDtypeKind, k2: IsDtypeKind):
-    return [
-        rewrite(isdtype(DType.float32, IsDtypeKind.string("integral"))).to(FALSE),
-        rewrite(isdtype(DType.float64, IsDtypeKind.string("integral"))).to(FALSE),
-        rewrite(isdtype(DType.object, IsDtypeKind.string("integral"))).to(FALSE),
-        rewrite(isdtype(DType.int64, IsDtypeKind.string("integral"))).to(TRUE),
-        rewrite(isdtype(DType.int32, IsDtypeKind.string("integral"))).to(TRUE),
-        rewrite(isdtype(DType.float32, IsDtypeKind.string("real floating"))).to(TRUE),
-        rewrite(isdtype(DType.float64, IsDtypeKind.string("real floating"))).to(TRUE),
-        rewrite(isdtype(DType.object, IsDtypeKind.string("real floating"))).to(FALSE),
-        rewrite(isdtype(DType.int64, IsDtypeKind.string("real floating"))).to(FALSE),
-        rewrite(isdtype(DType.int32, IsDtypeKind.string("real floating"))).to(FALSE),
-        rewrite(isdtype(DType.float32, IsDtypeKind.string("complex floating"))).to(FALSE),
-        rewrite(isdtype(DType.float64, IsDtypeKind.string("complex floating"))).to(FALSE),
-        rewrite(isdtype(DType.object, IsDtypeKind.string("complex floating"))).to(FALSE),
-        rewrite(isdtype(DType.int64, IsDtypeKind.string("complex floating"))).to(FALSE),
-        rewrite(isdtype(DType.int32, IsDtypeKind.string("complex floating"))).to(FALSE),
-        rewrite(isdtype(d, IsDtypeKind.NULL)).to(FALSE),
-        rewrite(isdtype(d, IsDtypeKind.dtype(d))).to(TRUE),
-        rewrite(isdtype(d, k1 | k2)).to(isdtype(d, k1) | isdtype(d, k2)),
-        rewrite(k1 | IsDtypeKind.NULL).to(k1),
+        rewrite(TRUE.if_int(i, j)).to(i),
+        rewrite(FALSE.if_int(i, j)).to(j),
     ]
 
 
@@ -264,10 +194,13 @@ converter(int, Int, lambda x: Int(x))
 
 
 class Float(Expr):
+    # Differentiate costs of three constructors so extraction is deterministic if all three are present
+    @method(cost=3)
     def __init__(self, value: f64Like) -> None: ...
 
     def abs(self) -> Float: ...
 
+    @method(cost=2)
     @classmethod
     def rational(cls, r: Rational) -> Float: ...
 
@@ -366,6 +299,85 @@ converter(type(None), OptionalInt, lambda _: OptionalInt.none)
 converter(Int, OptionalInt, OptionalInt.some)
 
 
+class DType(Expr):
+    float64: ClassVar[DType]
+    float32: ClassVar[DType]
+    int64: ClassVar[DType]
+    int32: ClassVar[DType]
+    object: ClassVar[DType]
+    bool: ClassVar[DType]
+
+    def __eq__(self, other: DType) -> Boolean:  # type: ignore[override]
+        ...
+
+
+float64 = DType.float64
+float32 = DType.float32
+int32 = DType.int32
+int64 = DType.int64
+
+_DTYPES = [float64, float32, int32, int64, DType.object]
+
+converter(type, DType, lambda x: convert(np.dtype(x), DType))
+converter(type(np.dtype), DType, lambda x: getattr(DType, x.name))  # type: ignore[call-overload]
+
+
+@array_api_ruleset.register
+def _():
+    for l, r in itertools.product(_DTYPES, repeat=2):
+        yield rewrite(l == r).to(TRUE if l is r else FALSE)
+
+
+class IsDtypeKind(Expr):
+    NULL: ClassVar[IsDtypeKind]
+
+    @classmethod
+    def string(cls, s: StringLike) -> IsDtypeKind: ...
+
+    @classmethod
+    def dtype(cls, d: DType) -> IsDtypeKind: ...
+
+    @method(cost=10)
+    def __or__(self, other: IsDtypeKind) -> IsDtypeKind: ...
+
+
+# TODO: Make kind more generic to support tuples.
+@function
+def isdtype(dtype: DType, kind: IsDtypeKind) -> Boolean: ...
+
+
+converter(DType, IsDtypeKind, lambda x: IsDtypeKind.dtype(x))
+converter(str, IsDtypeKind, lambda x: IsDtypeKind.string(x))
+converter(
+    tuple, IsDtypeKind, lambda x: convert(x[0], IsDtypeKind) | convert(x[1:], IsDtypeKind) if x else IsDtypeKind.NULL
+)
+
+
+@array_api_ruleset.register
+def _isdtype(d: DType, k1: IsDtypeKind, k2: IsDtypeKind):
+    return [
+        rewrite(isdtype(DType.float32, IsDtypeKind.string("integral"))).to(FALSE),
+        rewrite(isdtype(DType.float64, IsDtypeKind.string("integral"))).to(FALSE),
+        rewrite(isdtype(DType.object, IsDtypeKind.string("integral"))).to(FALSE),
+        rewrite(isdtype(DType.int64, IsDtypeKind.string("integral"))).to(TRUE),
+        rewrite(isdtype(DType.int32, IsDtypeKind.string("integral"))).to(TRUE),
+        rewrite(isdtype(DType.float32, IsDtypeKind.string("real floating"))).to(TRUE),
+        rewrite(isdtype(DType.float64, IsDtypeKind.string("real floating"))).to(TRUE),
+        rewrite(isdtype(DType.object, IsDtypeKind.string("real floating"))).to(FALSE),
+        rewrite(isdtype(DType.int64, IsDtypeKind.string("real floating"))).to(FALSE),
+        rewrite(isdtype(DType.int32, IsDtypeKind.string("real floating"))).to(FALSE),
+        rewrite(isdtype(DType.float32, IsDtypeKind.string("complex floating"))).to(FALSE),
+        rewrite(isdtype(DType.float64, IsDtypeKind.string("complex floating"))).to(FALSE),
+        rewrite(isdtype(DType.object, IsDtypeKind.string("complex floating"))).to(FALSE),
+        rewrite(isdtype(DType.int64, IsDtypeKind.string("complex floating"))).to(FALSE),
+        rewrite(isdtype(DType.int32, IsDtypeKind.string("complex floating"))).to(FALSE),
+        rewrite(isdtype(d, IsDtypeKind.NULL)).to(FALSE),
+        rewrite(isdtype(d, IsDtypeKind.dtype(d))).to(TRUE),
+        rewrite(isdtype(d, k1 | k2)).to(isdtype(d, k1) | isdtype(d, k2)),
+        rewrite(k1 | IsDtypeKind.NULL).to(k1),
+    ]
+
+
 class Slice(Expr):
     def __init__(
         self,

egglog/exp/array_api_numba.py

@@ -31,7 +31,12 @@ 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), subsume=True).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)

egglog/exp/siu_examples.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+import egglog
+
+from .array_api import Int
+
+# https://github.com/sklam/pyasir/blob/c363ff4f8f91177700ad4108dd5042b9b97d8289/pyasir/tests/test_fib.py
+
+# In progress - should be able to re-create this
+# @df.func
+# def fib_ir(n: pyasir.Int64) -> pyasir.Int64:
+#     @df.switch(n <= 1)
+#     def swt(n):
+#         @df.case(1)
+#         def case0(n):
+#             return 1
+
+#         @df.case(0)
+#         def case1(n):
+#             return fib_ir(n - 1) + fib_ir(n - 2)
+
+#         yield case0
+#         yield case1
+
+#     r = swt(n)
+#     return r
+
+
+# With something like this:
+@egglog.function
+def fib(n: Int) -> Int:
+    return (n <= Int(1)).if_int(
+        Int(1),
+        fib(n - Int(1)) + fib(n - Int(2)),
+    )