From f04704bf3ddd690e0d5dd2a911e24beb48d494ad Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Tue, 21 Apr 2026 17:20:43 -0400 Subject: [PATCH 01/26] add an op for dataclass construction --- effectful/ops/semantics.py | 32 ++++++++++++++++++++++---------- tests/test_ops_semantics.py | 21 +++++++++++++++++++++ 2 files changed, 43 insertions(+), 10 deletions(-) diff --git a/effectful/ops/semantics.py b/effectful/ops/semantics.py index f7678fd24..f38ea4cd7 100644 --- a/effectful/ops/semantics.py +++ b/effectful/ops/semantics.py @@ -211,19 +211,31 @@ def evaluate[T]( @evaluate.register(bytes) def _evaluate_object[T](expr: T, **kwargs) -> T: if dataclasses.is_dataclass(expr) and not isinstance(expr, type): - return typing.cast( - T, - dataclasses.replace( - expr, - **{ - field.name: evaluate(getattr(expr, field.name)) - for field in dataclasses.fields(expr) - }, - ), - ) + return _evaluate_dataclass(expr, **kwargs) return expr +def _get_dataclass_constr_op(typ): + if hasattr(typ, "constr_op"): + return typ.constr_op + + @Operation.define + def constr_op(*args, **kwargs) -> typ: + return typ(*args, **kwargs) + + typ.constr_op = constr_op + return constr_op + + +def _evaluate_dataclass[T](expr: T, **kwargs) -> T: + dataclass_op = _get_dataclass_constr_op(type(expr)) + subst = { + field.name: evaluate(getattr(expr, field.name)) + for field in dataclasses.fields(expr) + } + return typing.cast(T, dataclass_op(**subst)) + + @evaluate.register(Term) def _evaluate_term(expr: Term, **kwargs): args = tuple(evaluate(arg) for arg in expr.args) diff --git a/tests/test_ops_semantics.py b/tests/test_ops_semantics.py index 287c81769..85526c9c8 100644 --- a/tests/test_ops_semantics.py +++ b/tests/test_ops_semantics.py @@ -877,3 +877,24 @@ def __init__(self, x: int): v = Operation.define(int) assert fvsof(A(v())) == {v} + + +def test_defdata_dataclass_init_effects() -> None: + @Operation.define + def f(x: int): + raise NotHandled + + @dataclasses.dataclass + class A: + x: int + + def __init__(self, x: int): + self.x = f(x) + + @Operation.define + def g(a: A): + raise NotHandled + + v = Operation.define(int) + t = g(A(v())) + assert isinstance(t.args[0].x, Term) From a8bdc172067dfe81ca1fd07beaded5929efac71a Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Tue, 21 Apr 2026 17:25:51 -0400 Subject: [PATCH 02/26] lint --- effectful/ops/semantics.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/effectful/ops/semantics.py b/effectful/ops/semantics.py index f38ea4cd7..8059b538e 100644 --- a/effectful/ops/semantics.py +++ b/effectful/ops/semantics.py @@ -231,7 +231,7 @@ def _evaluate_dataclass[T](expr: T, **kwargs) -> T: dataclass_op = _get_dataclass_constr_op(type(expr)) subst = { field.name: evaluate(getattr(expr, field.name)) - for field in dataclasses.fields(expr) + for field in dataclasses.fields(expr) # type: ignore[arg-type] } return typing.cast(T, dataclass_op(**subst)) From 50c20caea4695ea8cc465f4dac7be5a65ad97af9 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Mon, 20 Jul 2026 14:10:25 -0400 Subject: [PATCH 03/26] add collection casts --- effectful/ops/semantics.py | 36 +++++++++++++++++++++++++++--------- effectful/ops/syntax.py | 36 +++++++++++++++++++++++++++++++++++- 2 files changed, 62 insertions(+), 10 deletions(-) diff --git a/effectful/ops/semantics.py b/effectful/ops/semantics.py index de041b61f..fb3727ee6 100644 --- a/effectful/ops/semantics.py +++ b/effectful/ops/semantics.py @@ -8,7 +8,14 @@ from collections.abc import Callable from typing import Any -from effectful.ops.syntax import _CustomSingleDispatchCallable, defdata, defop +from effectful.ops.syntax import ( + CollectionOperation, + _CustomSingleDispatchCallable, + as_list, + as_tuple, + defdata, + defop, +) from effectful.ops.types import ( Expr, Interpretation, @@ -115,6 +122,15 @@ def handler(intp: Interpretation): yield intp +@functools.cache +def _as_type(typ): + @CollectionOperation.define + def _as_typ(*args, **kwargs) -> typ: + return typ(*args, **kwargs) + + return _as_typ + + @_CustomSingleDispatchCallable def evaluate[T]( __dispatch: Callable[[type], Callable[..., Expr[T]]], @@ -183,17 +199,19 @@ def _evaluate_operation(expr: Operation, **kwargs) -> Operation: @evaluate.register(collections.defaultdict) def _evaluate_defaultdict(expr, **kwargs): - return type(expr)(expr.default_factory, evaluate(tuple(expr.items()))) + return _as_type(type(expr))( + expr.default_factory, as_list(*(evaluate(item) for item in expr.items())) + ) @evaluate.register(types.MappingProxyType) def _evaluate_mappingproxytype(expr, **kwargs): - return type(expr)(dict(evaluate(tuple(expr.items())))) + return _as_type(type(expr))(as_list(*(evaluate(item) for item in expr.items()))) @evaluate.register(collections.abc.Mapping) def _evaluate_mapping(expr, **kwargs): - return type(expr)(evaluate(tuple(expr.items()))) + return _as_type(type(expr))(as_list(*(evaluate(item) for item in expr.items()))) @evaluate.register(tuple) @@ -203,27 +221,27 @@ def _evaluate_tuple(expr, **kwargs): and hasattr(expr, "_fields") and all(hasattr(expr, field) for field in getattr(expr, "_fields")) ): # namedtuple - return type(expr)( + return _as_type(type(expr))( **{field: evaluate(getattr(expr, field)) for field in expr._fields} ) else: - return type(expr)(evaluate(item) for item in expr) + return _as_type(type(expr))(as_tuple(*(evaluate(item) for item in expr))) @evaluate.register(collections.abc.Sequence) def _evaluate_sequence(expr, **kwargs): - return type(expr)(evaluate(item) for item in expr) + return _as_type(type(expr))(as_list(*(evaluate(item) for item in expr))) @evaluate.register(collections.abc.ItemsView) @evaluate.register(collections.abc.KeysView) def _evaluate_set_view(expr, **kwargs): - return {evaluate(item) for item in expr} + return as_set(*(evaluate(item) for item in expr)) @evaluate.register(collections.abc.ValuesView) def _evaluate_list_view(expr, **kwargs): - return [evaluate(item) for item in expr] + return as_list(*(evaluate(item) for item in expr)) def _simple_type(tp: type) -> type: diff --git a/effectful/ops/syntax.py b/effectful/ops/syntax.py index 764016752..817241623 100644 --- a/effectful/ops/syntax.py +++ b/effectful/ops/syntax.py @@ -514,7 +514,18 @@ def apply_cast(op, *args, **kwargs): dispatch_type = _simple_type(full_type.value) return __dispatch(dispatch_type)(dispatch_type, op, *args, **kwargs) - analysis = productN({typ: {apply: apply_type}, cast: {apply: apply_cast}}) + analysis = productN( + { + typ: { + apply: apply_type, + CollectionOperation.__apply__: apply.__default_rule__, + }, + cast: { + apply: apply_cast, + CollectionOperation.__apply__: apply.__default_rule__, + }, + } + ) def evaluate_with_renaming(expr, ctx): """Evaluate an expression with renaming applied.""" @@ -1327,3 +1338,26 @@ class _IntegralTerm[T: numbers.Integral](_RationalTerm[T]): @defdata.register(bool) class _BoolTerm[T: bool](_IntegralTerm[T]): # type: ignore pass + + +class CollectionOperation(Operation): ... + + +@CollectionOperation.define +def as_tuple(*args) -> tuple: + return tuple(args) + + +@CollectionOperation.define +def as_list[T](*args: T) -> list[T]: + return list(args) + + +@CollectionOperation.define +def as_set[T](*args: T) -> set[T]: + return set(args) + + +@CollectionOperation.define +def as_dict[K, V](*args: tuple[K, V]) -> dict[K, V]: + return dict(args) From cc5cc62ab588a4e5e792f80d3c521c0936f1a569 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Mon, 20 Jul 2026 14:11:25 -0400 Subject: [PATCH 04/26] use shared code --- effectful/ops/semantics.py | 14 +------------- 1 file changed, 1 insertion(+), 13 deletions(-) diff --git a/effectful/ops/semantics.py b/effectful/ops/semantics.py index 1377bf477..f2ad03a96 100644 --- a/effectful/ops/semantics.py +++ b/effectful/ops/semantics.py @@ -173,20 +173,8 @@ def _evaluate_object[T](expr: T, **kwargs) -> T: return expr -def _get_dataclass_constr_op(typ): - if hasattr(typ, "constr_op"): - return typ.constr_op - - @Operation.define - def constr_op(*args, **kwargs) -> typ: - return typ(*args, **kwargs) - - typ.constr_op = constr_op - return constr_op - - def _evaluate_dataclass[T](expr: T, **kwargs) -> T: - dataclass_op = _get_dataclass_constr_op(type(expr)) + dataclass_op = _as_type(type(expr)) subst = { field.name: evaluate(getattr(expr, field.name)) for field in dataclasses.fields(expr) # type: ignore[arg-type] From c31ceda53977a5b807cc98bce3c9090fe7e6f767 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Mon, 20 Jul 2026 14:50:25 -0400 Subject: [PATCH 05/26] wip --- effectful/ops/semantics.py | 45 ++++++++++++++++++++++++-------------- effectful/ops/syntax.py | 14 ++++++------ 2 files changed, 35 insertions(+), 24 deletions(-) diff --git a/effectful/ops/semantics.py b/effectful/ops/semantics.py index f2ad03a96..1686493e7 100644 --- a/effectful/ops/semantics.py +++ b/effectful/ops/semantics.py @@ -9,9 +9,10 @@ from typing import Any from effectful.ops.syntax import ( - CollectionOperation, + CollectionConstrOperation, _CustomSingleDispatchCallable, as_list, + as_set, as_tuple, defdata, defop, @@ -122,9 +123,12 @@ def handler(intp: Interpretation): yield intp +class DataclassConstrOperation(Operation): ... + + @functools.cache -def _as_type(typ): - @CollectionOperation.define +def _as_type(typ, operation_type=CollectionConstrOperation): + @operation_type.define def _as_typ(*args, **kwargs) -> typ: return typ(*args, **kwargs) @@ -174,7 +178,7 @@ def _evaluate_object[T](expr: T, **kwargs) -> T: def _evaluate_dataclass[T](expr: T, **kwargs) -> T: - dataclass_op = _as_type(type(expr)) + dataclass_op = _as_type(type(expr), operation_type=DataclassConstrOperation) subst = { field.name: evaluate(getattr(expr, field.name)) for field in dataclasses.fields(expr) # type: ignore[arg-type] @@ -294,7 +298,9 @@ def typeof[T](term: Expr[T]) -> type[T]: def _apply(op, *args, **kwargs): return Box(op.__type_rule__(*args, **kwargs)) - with interpreter({apply: _apply}): + with interpreter( + {apply: _apply, CollectionConstrOperation.__apply__: apply.__default_rule__} + ): type_or_value = evaluate(term) if isinstance(type_or_value, Box): return _simple_type(type_or_value.value) @@ -315,18 +321,23 @@ def fvsof[S](term: Expr[S]) -> collections.abc.Set[Operation]: """ from effectful.internals.runtime import interpreter - _fvs: set[Operation] = set() - - def _update_fvs(op, *args, **kwargs): - _fvs.add(op) + def _apply(op, *args, **kwargs): + free_vars = set().union( + {op}, + *( + {x} if isinstance(x, Operation) else x + for x in (*args, *kwargs.values()) + ), + ) bindings = op.__fvs_rule__(*args, **kwargs) - for bound_var in set().union(*(*bindings.args, *bindings.kwargs.values())): - assert isinstance(bound_var, Operation) - if bound_var in _fvs: - _fvs.remove(bound_var) - return defdata(op, *args, **kwargs) + bound_vars = set().union(*bindings.args, *bindings.kwargs.values()) + return free_vars - bound_vars - with interpreter({apply: _update_fvs}): - evaluate(term) + with interpreter({apply: _apply}): + fvs = evaluate(term) - return _fvs + return { + op + for op in fvs + if not isinstance(op, DataclassConstrOperation | CollectionConstrOperation) + } diff --git a/effectful/ops/syntax.py b/effectful/ops/syntax.py index 817241623..2f5a53d2d 100644 --- a/effectful/ops/syntax.py +++ b/effectful/ops/syntax.py @@ -518,11 +518,11 @@ def apply_cast(op, *args, **kwargs): { typ: { apply: apply_type, - CollectionOperation.__apply__: apply.__default_rule__, + CollectionConstrOperation.__apply__: apply.__default_rule__, }, cast: { apply: apply_cast, - CollectionOperation.__apply__: apply.__default_rule__, + CollectionConstrOperation.__apply__: apply.__default_rule__, }, } ) @@ -1340,24 +1340,24 @@ class _BoolTerm[T: bool](_IntegralTerm[T]): # type: ignore pass -class CollectionOperation(Operation): ... +class CollectionConstrOperation(Operation): ... -@CollectionOperation.define +@CollectionConstrOperation.define def as_tuple(*args) -> tuple: return tuple(args) -@CollectionOperation.define +@CollectionConstrOperation.define def as_list[T](*args: T) -> list[T]: return list(args) -@CollectionOperation.define +@CollectionConstrOperation.define def as_set[T](*args: T) -> set[T]: return set(args) -@CollectionOperation.define +@CollectionConstrOperation.define def as_dict[K, V](*args: tuple[K, V]) -> dict[K, V]: return dict(args) From 17b40ec1c4645294016f09891b8537d2f81dc94f Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Tue, 21 Jul 2026 18:24:56 -0400 Subject: [PATCH 06/26] rework fvsof to prepare for caching --- effectful/ops/semantics.py | 105 +++++++++++++++++++++++++++++++----- effectful/ops/types.py | 11 +++- tests/test_ops_semantics.py | 56 +++++++++++++++---- 3 files changed, 146 insertions(+), 26 deletions(-) diff --git a/effectful/ops/semantics.py b/effectful/ops/semantics.py index 1686493e7..95175fc9e 100644 --- a/effectful/ops/semantics.py +++ b/effectful/ops/semantics.py @@ -10,6 +10,7 @@ from effectful.ops.syntax import ( CollectionConstrOperation, + _BaseTerm, _CustomSingleDispatchCallable, as_list, as_set, @@ -234,7 +235,9 @@ def _evaluate_tuple(expr, **kwargs): @evaluate.register(collections.abc.Sequence) def _evaluate_sequence(expr, **kwargs): - return _as_type(type(expr))(as_list(*(evaluate(item) for item in expr))) + seq = as_list(*(evaluate(item) for item in expr)) + cast = _as_type(type(expr)) + return cast(seq) @evaluate.register(collections.abc.ItemsView) @@ -307,6 +310,40 @@ def _apply(op, *args, **kwargs): return typing.cast(type[T], type(type_or_value)) +# def _binders_apply(op, *args, **kwargs): +# return frozenset() + + +# def _binders_collection_apply(op, *args, **kwargs): +# return frozenset().union( +# *( +# x +# if isinstance(x, frozenset) +# else {x} +# if isinstance(x, Operation) +# else set() +# for x in args +# ) +# ) + + +# _BINDERS_INTP = { +# apply: _binders_apply, +# CollectionConstrOperation.__apply__: _binders_collection_apply, +# } + + +# def binders(term: Expr[S]) -> collections.abc.Set[Operation]: +# with interpreter(_BINDERS_INTP): +# bs = evaluate(term) +# if isinstance(bs, frozenset): +# return bs +# elif isinstance(bs, Operation): +# return {bs} +# else: +# return frozenset() + + def fvsof[S](term: Expr[S]) -> collections.abc.Set[Operation]: """Return the free variables of an expression. @@ -319,25 +356,65 @@ def fvsof[S](term: Expr[S]) -> collections.abc.Set[Operation]: >>> assert f in fvs >>> assert len(fvs) == 1 """ + from effectful.internals.product_n import _unpack, productN from effectful.internals.runtime import interpreter - def _apply(op, *args, **kwargs): - free_vars = set().union( + # Analysis for type computation and term reconstruction + _fvsof_fvs = defop(object, name="fvsof_fvs") + _fvsof_binders = defop(object, name="fvsof_binders") + + def _apply_collection_binders(op, *args, **kwargs): + return frozenset().union( + *( + {x} + if isinstance(x, Operation) + else x + if isinstance(x, frozenset) + else set() + for x in (*args, *kwargs.values()) + ) + ) + + def _apply_binders(op, *args, **kwargs): + args = tuple(frozenset() if isinstance(x, Term) else x for x in args) + kwargs = { + k: frozenset() if isinstance(v, Term) else v for (k, v) in kwargs.items() + } + return _BaseTerm(op, *args, **kwargs) + + def _apply_fvs(op, *args, **kwargs): + term = _fvsof_binders() + + if isinstance(term, Term): + bindings = op.__fvs_rule__(*term.args, **term.kwargs) + binders = frozenset().union(*(*bindings.args, *bindings.kwargs.values())) + else: + binders = frozenset() + + fvs = frozenset().union( {op}, *( - {x} if isinstance(x, Operation) else x + x if isinstance(x, frozenset) else frozenset() for x in (*args, *kwargs.values()) ), ) - bindings = op.__fvs_rule__(*args, **kwargs) - bound_vars = set().union(*bindings.args, *bindings.kwargs.values()) - return free_vars - bound_vars + fvs -= binders + return fvs + + _fvsof_intp = productN( + { + _fvsof_fvs: {apply: _apply_fvs}, + _fvsof_binders: { + CollectionConstrOperation.__apply__: _apply_collection_binders, + apply: _apply_binders, + }, + } + ) - with interpreter({apply: _apply}): - fvs = evaluate(term) + with interpreter(_fvsof_intp): + result = evaluate(term) - return { - op - for op in fvs - if not isinstance(op, DataclassConstrOperation | CollectionConstrOperation) - } + fvs = _unpack(result, _fvsof_fvs) + if not isinstance(fvs, frozenset): + return frozenset() + return fvs diff --git a/effectful/ops/types.py b/effectful/ops/types.py index 65e31de38..625be7ca2 100644 --- a/effectful/ops/types.py +++ b/effectful/ops/types.py @@ -465,11 +465,18 @@ def __set_name__[T](self, owner: type[T], name: str) -> None: def __get__[T](self, instance: T | None, owner: type[T] | None = None): if hasattr(instance, "__dict__") and hasattr(self, "_name_on_instance"): - from effectful.ops.semantics import fvsof + from effectful.ops.semantics import ( + DataclassConstrOperation, + _as_type, + fvsof, + ) if self._name_on_instance in instance.__dict__: return instance.__dict__[self._name_on_instance] - elif isinstance(instance, Term) or fvsof(instance): + elif isinstance(instance, Term) or ( + fvsof(instance) + - {_as_type(owner, operation_type=DataclassConstrOperation)} + ): return types.MethodType(self, instance) else: diff --git a/tests/test_ops_semantics.py b/tests/test_ops_semantics.py index 1c2e548ce..bdccb139b 100644 --- a/tests/test_ops_semantics.py +++ b/tests/test_ops_semantics.py @@ -457,18 +457,27 @@ def Nested(*args, **kwargs): assert evaluate(t) == Nested([{"a": 2}, 1, (1, 2)], 1, arg1={"b": 1}) -def test_ctxof(): - x = defop(object) - y = defop(object) +@pytest.mark.parametrize( + "build_args", + [ + lambda x, y: (x(), y()), + lambda x, y: ([x()], y()), + lambda x, y: ([x()], [y()]), + lambda x, y: (([x()], [y()]),), + ], +) +def test_ctxof(build_args): + x = defop(object, name="x") + y = defop(object, name="y") @defop def Nested(*args, **kwargs): raise NotHandled - assert fvsof(Nested(x(), y())) >= {x, y} - assert fvsof(Nested([x()], y())) >= {x, y} - assert fvsof(Nested([x()], [y()])) >= {x, y} - assert fvsof(Nested((x(), y()))) >= {x, y} + term = Nested(*build_args(x, y)) + actual = fvsof(term) + expected = {x, y, Nested} + assert actual >= expected def test_handler_typing() -> None: @@ -722,8 +731,9 @@ def Lam2[A, B]( raise NotHandled term = Lam2(add(x(), add(y(), z())), x, y) - assert not {x, y} <= fvsof(term) - assert fvsof(term) == {z, Lam2, add} + actual = fvsof(term) + assert not ({x, y} & actual) + assert actual >= {z, Lam2, add} def test_interpretation_typing(): @@ -784,7 +794,8 @@ def __init__(self, x: int): self.x = x v = Operation.define(int) - assert fvsof(A(v())) == {v} + actual = fvsof(A(v())) + assert actual >= {v} def test_defdata_dataclass_init_effects() -> None: @@ -839,6 +850,31 @@ def f(self): assert b.f() == "*B*" +def test_instanceop_dataclass() -> None: + """Dataclasses with no free variables get instance operations.""" + + @dataclasses.dataclass + class A: + @Operation.define + def f(self): + raise NotHandled + + assert isinstance(A.f, Operation) + assert isinstance(A().f, Operation) + + @dataclasses.dataclass + class B: + x: int + + @Operation.define + def g(self): + raise NotHandled + + assert isinstance(B.g, Operation) + fv = Operation.define(int)() + assert not isinstance(B(fv).g, Operation) + + def test_coproduct_fwd_chain(benchmark): """Benchmark coproduct + fwd over a deep chain of forwarding handlers. From e2cb168b2fb55e95884a84b256bd889b23db027b Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Wed, 22 Jul 2026 11:05:19 -0400 Subject: [PATCH 07/26] start translating sizesof --- effectful/handlers/jax/_handlers.py | 87 +++++++++++++++++------- effectful/handlers/jax/numpy/__init__.py | 28 ++++++-- effectful/ops/semantics.py | 35 ---------- 3 files changed, 83 insertions(+), 67 deletions(-) diff --git a/effectful/handlers/jax/_handlers.py b/effectful/handlers/jax/_handlers.py index 308cdb76e..2fbc92885 100644 --- a/effectful/handlers/jax/_handlers.py +++ b/effectful/handlers/jax/_handlers.py @@ -1,4 +1,5 @@ import functools +import itertools import typing from collections.abc import Callable, Mapping, Sequence from types import EllipsisType @@ -13,7 +14,9 @@ from effectful.internals.runtime import interpreter from effectful.ops.semantics import apply, evaluate, fvsof, typeof from effectful.ops.syntax import ( + CollectionConstrOperation, Scoped, + _BaseTerm, _CustomSingleDispatchCallable, defdata, deffn, @@ -39,7 +42,7 @@ def is_eager_array(x): ) -def sizesof(value) -> Mapping[Operation[[], jax.Array], int]: +def sizesof(term: Expr) -> Mapping[Operation[[], jax.Array], int]: """Return the sizes of named dimensions in an array expression. Sizes are inferred from the array shape. @@ -53,30 +56,64 @@ def sizesof(value) -> Mapping[Operation[[], jax.Array], int]: >>> sizes = sizesof(jax_getitem(jnp.ones((2, 3)), [a(), b()])) >>> assert sizes[a] == 2 and sizes[b] == 3 """ - sizes: dict[Operation[[], jax.Array], int] = {} - - def update_sizes(sizes, op, size): - old_size = sizes.get(op) - if old_size is not None and size != old_size: - raise ValueError( - f"Named index {op} used in incompatible dimensions of size {old_size} and {size}" - ) - sizes[op] = size - - def _getitem_sizeof(x: jax.Array, key: tuple[Expr[IndexElement], ...]): - if is_eager_array(x): - for i, k in enumerate(key): - if isinstance(k, Term) and len(k.args) == 0 and len(k.kwargs) == 0: - update_sizes(sizes, k.op, x.shape[i]) - return defdata(jax_getitem, x, key) - - def _apply(op, *args, **kwargs): - return defdata(op, *args, **kwargs) - - with interpreter({jax_getitem: _getitem_sizeof, apply: _apply}): - evaluate(value) - - return sizes + from effectful.internals.product_n import _unpack, productN + + # Analysis for type computation and term reconstruction + _sizes = defop(object, name="sizes") + _getitem_term = defop(object, name="getitem_args") + + def _retain(op, *args, **kwargs): + return _BaseTerm(op, *args, **kwargs) + + def _merge(s1, s2): + s3 = s1.copy() + for k, v in s2.items(): + if k in s3 and s3[k] != v: + raise ValueError( + f"Named index {k} used in incompatible dimensions of size {s3[k]} and {v}" + ) + s3[k] = v + return s3 + + def _apply_sizes(op, *args, **kwargs): + analyses = ( + x for x in (*args, *kwargs.values()) if isinstance(x, dict) + ) + return functools.reduce(_merge, analyses, {}) + + def _getitem(arr, index): + term = _getitem_term() + assert isinstance(term, Term) + term_arr, term_index = term.args + + arg_sizes = (x for x in (arr, index) if isinstance(x, dict)) + if isinstance(term_arr, Term): + return functools.reduce(_merge, arg_sizes, {}) + + sizes = ( + {k.op: term_arr.shape[i]} + for i, k in enumerate(term_index) + if isinstance(k, Term) and len(k.args) == 0 and len(k.kwargs) == 0 + ) + return functools.reduce(_merge, itertools.chain(arg_sizes, sizes), {}) + + _intp = productN( + { + _sizes: {apply: _apply_sizes, jax_getitem: _getitem}, + _getitem_term: { + apply: _retain, + CollectionConstrOperation.__apply__: apply.__default_rule__, + }, + } + ) + + with interpreter(_intp): + result = evaluate(term) + + fvs = _unpack(result, _sizes) + if not isinstance(fvs, dict): + return {} + return fvs def _partial_eval(t: Expr[jax.Array]) -> Expr[jax.Array]: diff --git a/effectful/handlers/jax/numpy/__init__.py b/effectful/handlers/jax/numpy/__init__.py index 990830d27..d39f46c4e 100644 --- a/effectful/handlers/jax/numpy/__init__.py +++ b/effectful/handlers/jax/numpy/__init__.py @@ -1,3 +1,4 @@ +import types from typing import TYPE_CHECKING import jax.numpy @@ -7,15 +8,28 @@ _no_overload = ["array", "asarray"] for name, op in jax.numpy.__dict__.items(): - if not callable(op): + if isinstance(op, types.ModuleType): continue - jax_op = ( - _register_jax_op_no_partial_eval(op) - if name in _no_overload - else _register_jax_op(op) - ) - globals()[name] = jax_op + # copy constants + if isinstance(op, float | types.NoneType): + globals()[name] = op + + if callable(op): + if name == "__getattr__": + continue + + elif name in _no_overload: + globals()[name] = _register_jax_op_no_partial_eval(op) + + else: + globals()[name] = _register_jax_op(op) + jax_op = ( + _register_jax_op_no_partial_eval(op) + if name in _no_overload + else _register_jax_op(op) + ) + globals()[name] = jax_op pi = jax.numpy.pi diff --git a/effectful/ops/semantics.py b/effectful/ops/semantics.py index 95175fc9e..e1253b541 100644 --- a/effectful/ops/semantics.py +++ b/effectful/ops/semantics.py @@ -15,7 +15,6 @@ as_list, as_set, as_tuple, - defdata, defop, ) from effectful.ops.types import ( @@ -310,40 +309,6 @@ def _apply(op, *args, **kwargs): return typing.cast(type[T], type(type_or_value)) -# def _binders_apply(op, *args, **kwargs): -# return frozenset() - - -# def _binders_collection_apply(op, *args, **kwargs): -# return frozenset().union( -# *( -# x -# if isinstance(x, frozenset) -# else {x} -# if isinstance(x, Operation) -# else set() -# for x in args -# ) -# ) - - -# _BINDERS_INTP = { -# apply: _binders_apply, -# CollectionConstrOperation.__apply__: _binders_collection_apply, -# } - - -# def binders(term: Expr[S]) -> collections.abc.Set[Operation]: -# with interpreter(_BINDERS_INTP): -# bs = evaluate(term) -# if isinstance(bs, frozenset): -# return bs -# elif isinstance(bs, Operation): -# return {bs} -# else: -# return frozenset() - - def fvsof[S](term: Expr[S]) -> collections.abc.Set[Operation]: """Return the free variables of an expression. From 71f3283e5dc0c5ec05c94a228608abacfd720e71 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Wed, 22 Jul 2026 11:11:19 -0400 Subject: [PATCH 08/26] fix bug in fvsof --- effectful/handlers/jax/_handlers.py | 4 +--- effectful/ops/semantics.py | 14 +++++++++++++- tests/test_ops_semantics.py | 6 ++++++ 3 files changed, 20 insertions(+), 4 deletions(-) diff --git a/effectful/handlers/jax/_handlers.py b/effectful/handlers/jax/_handlers.py index 2fbc92885..c7195ff1f 100644 --- a/effectful/handlers/jax/_handlers.py +++ b/effectful/handlers/jax/_handlers.py @@ -76,9 +76,7 @@ def _merge(s1, s2): return s3 def _apply_sizes(op, *args, **kwargs): - analyses = ( - x for x in (*args, *kwargs.values()) if isinstance(x, dict) - ) + analyses = (x for x in (*args, *kwargs.values()) if isinstance(x, dict)) return functools.reduce(_merge, analyses, {}) def _getitem(arr, index): diff --git a/effectful/ops/semantics.py b/effectful/ops/semantics.py index e1253b541..bf420f34e 100644 --- a/effectful/ops/semantics.py +++ b/effectful/ops/semantics.py @@ -347,6 +347,15 @@ def _apply_binders(op, *args, **kwargs): } return _BaseTerm(op, *args, **kwargs) + def _apply_collection_fvs(op, *args, **kwargs): + return frozenset().union( + *( + x + for x in (*args, *kwargs.values()) + if isinstance(x, frozenset) + ) + ) + def _apply_fvs(op, *args, **kwargs): term = _fvsof_binders() @@ -368,7 +377,10 @@ def _apply_fvs(op, *args, **kwargs): _fvsof_intp = productN( { - _fvsof_fvs: {apply: _apply_fvs}, + _fvsof_fvs: { + apply: _apply_fvs, + CollectionConstrOperation.__apply__: _apply_collection_fvs, + }, _fvsof_binders: { CollectionConstrOperation.__apply__: _apply_collection_binders, apply: _apply_binders, diff --git a/tests/test_ops_semantics.py b/tests/test_ops_semantics.py index bdccb139b..fdc6f1cd7 100644 --- a/tests/test_ops_semantics.py +++ b/tests/test_ops_semantics.py @@ -736,6 +736,12 @@ def Lam2[A, B]( assert actual >= {z, Lam2, add} +def test_fvsof_collection_does_not_include_apply(): + x = defop(int, name="x") + + assert fvsof((x(),)) == {x} + + def test_interpretation_typing(): @defop def f[T](m: Mapping[Operation, T], x: T) -> T: From 68a84638d1741de5673e21af7fb914985ccd9ba8 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Wed, 22 Jul 2026 11:20:00 -0400 Subject: [PATCH 09/26] fix remaining tests --- effectful/handlers/jax/_handlers.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/effectful/handlers/jax/_handlers.py b/effectful/handlers/jax/_handlers.py index c7195ff1f..d09b801e2 100644 --- a/effectful/handlers/jax/_handlers.py +++ b/effectful/handlers/jax/_handlers.py @@ -63,7 +63,10 @@ def sizesof(term: Expr) -> Mapping[Operation[[], jax.Array], int]: _getitem_term = defop(object, name="getitem_args") def _retain(op, *args, **kwargs): - return _BaseTerm(op, *args, **kwargs) + return defdata(op, *args, **kwargs) + + def _retain_getitem(*args, **kwargs): + return defdata(jax_getitem, *args, **kwargs) def _merge(s1, s2): s3 = s1.copy() @@ -85,7 +88,7 @@ def _getitem(arr, index): term_arr, term_index = term.args arg_sizes = (x for x in (arr, index) if isinstance(x, dict)) - if isinstance(term_arr, Term): + if not is_eager_array(term_arr): return functools.reduce(_merge, arg_sizes, {}) sizes = ( @@ -100,6 +103,7 @@ def _getitem(arr, index): _sizes: {apply: _apply_sizes, jax_getitem: _getitem}, _getitem_term: { apply: _retain, + jax_getitem: _retain_getitem, CollectionConstrOperation.__apply__: apply.__default_rule__, }, } From 08fc2ceb2ca7abdce80479437fd835d617ae2e84 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Wed, 22 Jul 2026 11:23:09 -0400 Subject: [PATCH 10/26] simplify analysis --- effectful/handlers/jax/_handlers.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/effectful/handlers/jax/_handlers.py b/effectful/handlers/jax/_handlers.py index d09b801e2..bb3ad8b86 100644 --- a/effectful/handlers/jax/_handlers.py +++ b/effectful/handlers/jax/_handlers.py @@ -56,14 +56,16 @@ def sizesof(term: Expr) -> Mapping[Operation[[], jax.Array], int]: >>> sizes = sizesof(jax_getitem(jnp.ones((2, 3)), [a(), b()])) >>> assert sizes[a] == 2 and sizes[b] == 3 """ - from effectful.internals.product_n import _unpack, productN + from effectful.internals.product_n import _unpack, argsof, productN # Analysis for type computation and term reconstruction _sizes = defop(object, name="sizes") _getitem_term = defop(object, name="getitem_args") def _retain(op, *args, **kwargs): - return defdata(op, *args, **kwargs) + # Non-getitem subterms are opaque to this analysis. Keeping their + # arguments would retain the entire input term unnecessarily. + return _BaseTerm(op) def _retain_getitem(*args, **kwargs): return defdata(jax_getitem, *args, **kwargs) @@ -83,9 +85,10 @@ def _apply_sizes(op, *args, **kwargs): return functools.reduce(_merge, analyses, {}) def _getitem(arr, index): - term = _getitem_term() - assert isinstance(term, Term) - term_arr, term_index = term.args + # Inspect this getitem's arguments in the term projection without + # forcing that projection to retain the getitem result. + term_args, _ = argsof(_getitem_term) + term_arr, term_index = term_args arg_sizes = (x for x in (arr, index) if isinstance(x, dict)) if not is_eager_array(term_arr): From dac028eca3bdac4178ac372e91a64437e5bd92e1 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Wed, 22 Jul 2026 11:26:13 -0400 Subject: [PATCH 11/26] shrink terms --- effectful/ops/semantics.py | 30 +++++++++++++++++------------- 1 file changed, 17 insertions(+), 13 deletions(-) diff --git a/effectful/ops/semantics.py b/effectful/ops/semantics.py index bf420f34e..6be7ab47f 100644 --- a/effectful/ops/semantics.py +++ b/effectful/ops/semantics.py @@ -321,7 +321,7 @@ def fvsof[S](term: Expr[S]) -> collections.abc.Set[Operation]: >>> assert f in fvs >>> assert len(fvs) == 1 """ - from effectful.internals.product_n import _unpack, productN + from effectful.internals.product_n import _unpack, argsof, productN from effectful.internals.runtime import interpreter # Analysis for type computation and term reconstruction @@ -341,11 +341,9 @@ def _apply_collection_binders(op, *args, **kwargs): ) def _apply_binders(op, *args, **kwargs): - args = tuple(frozenset() if isinstance(x, Term) else x for x in args) - kwargs = { - k: frozenset() if isinstance(v, Term) else v for (k, v) in kwargs.items() - } - return _BaseTerm(op, *args, **kwargs) + # Parent operations only need to know that this child is a term. Its + # arguments are available through argsof while this node is analyzed. + return _BaseTerm(op) def _apply_collection_fvs(op, *args, **kwargs): return frozenset().union( @@ -357,13 +355,19 @@ def _apply_collection_fvs(op, *args, **kwargs): ) def _apply_fvs(op, *args, **kwargs): - term = _fvsof_binders() - - if isinstance(term, Term): - bindings = op.__fvs_rule__(*term.args, **term.kwargs) - binders = frozenset().union(*(*bindings.args, *bindings.kwargs.values())) - else: - binders = frozenset() + binder_args, binder_kwargs = argsof(_fvsof_binders) + # This rule handles Operation.__apply__ directly, so its first argument + # is the operation being applied rather than an argument to that + # operation. + binder_args = tuple( + frozenset() if isinstance(x, Term) else x for x in binder_args[1:] + ) + binder_kwargs = { + k: frozenset() if isinstance(v, Term) else v + for k, v in binder_kwargs.items() + } + bindings = op.__fvs_rule__(*binder_args, **binder_kwargs) + binders = frozenset().union(*(*bindings.args, *bindings.kwargs.values())) fvs = frozenset().union( {op}, From 4bd9bd976bafe29a80c36e21d5a139d30616e0c9 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Wed, 22 Jul 2026 11:31:15 -0400 Subject: [PATCH 12/26] redundant --- effectful/handlers/jax/numpy/__init__.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/effectful/handlers/jax/numpy/__init__.py b/effectful/handlers/jax/numpy/__init__.py index d39f46c4e..cc20d7498 100644 --- a/effectful/handlers/jax/numpy/__init__.py +++ b/effectful/handlers/jax/numpy/__init__.py @@ -31,8 +31,6 @@ ) globals()[name] = jax_op -pi = jax.numpy.pi - # Tell mypy about our wrapped functions. if TYPE_CHECKING: from jax.numpy import * # noqa: F403 From c371bf19ece02236f0aae24080068d4135565683 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Wed, 22 Jul 2026 11:42:56 -0400 Subject: [PATCH 13/26] lint --- effectful/ops/semantics.py | 13 +++++-------- effectful/ops/types.py | 3 ++- 2 files changed, 7 insertions(+), 9 deletions(-) diff --git a/effectful/ops/semantics.py b/effectful/ops/semantics.py index 6be7ab47f..85ad2a72c 100644 --- a/effectful/ops/semantics.py +++ b/effectful/ops/semantics.py @@ -127,9 +127,9 @@ class DataclassConstrOperation(Operation): ... @functools.cache -def _as_type(typ, operation_type=CollectionConstrOperation): +def _as_type(typ: type, operation_type=CollectionConstrOperation): @operation_type.define - def _as_typ(*args, **kwargs) -> typ: + def _as_typ(*args, **kwargs) -> typ: # type: ignore[valid-type] return typ(*args, **kwargs) return _as_typ @@ -178,7 +178,8 @@ def _evaluate_object[T](expr: T, **kwargs) -> T: def _evaluate_dataclass[T](expr: T, **kwargs) -> T: - dataclass_op = _as_type(type(expr), operation_type=DataclassConstrOperation) + typ: type = type(expr) + dataclass_op = _as_type(typ, operation_type=DataclassConstrOperation) subst = { field.name: evaluate(getattr(expr, field.name)) for field in dataclasses.fields(expr) # type: ignore[arg-type] @@ -347,11 +348,7 @@ def _apply_binders(op, *args, **kwargs): def _apply_collection_fvs(op, *args, **kwargs): return frozenset().union( - *( - x - for x in (*args, *kwargs.values()) - if isinstance(x, frozenset) - ) + *(x for x in (*args, *kwargs.values()) if isinstance(x, frozenset)) ) def _apply_fvs(op, *args, **kwargs): diff --git a/effectful/ops/types.py b/effectful/ops/types.py index 625be7ca2..106c6c646 100644 --- a/effectful/ops/types.py +++ b/effectful/ops/types.py @@ -471,11 +471,12 @@ def __get__[T](self, instance: T | None, owner: type[T] | None = None): fvsof, ) + instance_type: type = type(instance) if self._name_on_instance in instance.__dict__: return instance.__dict__[self._name_on_instance] elif isinstance(instance, Term) or ( fvsof(instance) - - {_as_type(owner, operation_type=DataclassConstrOperation)} + - {_as_type(instance_type, operation_type=DataclassConstrOperation)} ): return types.MethodType(self, instance) else: From 4e8f9bea64baf8d14ff508f7df383b24e8f189b7 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Thu, 23 Jul 2026 09:25:53 -0400 Subject: [PATCH 14/26] update fvsof documentation --- effectful/ops/semantics.py | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/effectful/ops/semantics.py b/effectful/ops/semantics.py index 85ad2a72c..52134c771 100644 --- a/effectful/ops/semantics.py +++ b/effectful/ops/semantics.py @@ -311,16 +311,30 @@ def _apply(op, *args, **kwargs): def fvsof[S](term: Expr[S]) -> collections.abc.Set[Operation]: - """Return the free variables of an expression. + """Return the free operations in a term. + + An operation belongs to `fvsof(t)` when it appears free in the term `t`. + This excludes operations like `apply` or collection constructors that are + raised during `evaluate` but do not appear in `t`. It also excludes + operations that are bound by a `Scoped` operation. However, it is not + restricted to the nullary operations in `t`. **Example usage**: + `fvsof` includes all unbound operations in a term: + + >>> a = defop(int) >>> @defop ... def f(x: int, y: int) -> int: ... raise NotHandled - >>> fvs = fvsof(f(1, 2)) - >>> assert f in fvs - >>> assert len(fvs) == 1 + >>> fvs = fvsof(f(a(), 1)) + >>> assert fvs >= {f, a} + + `fvsof` accepts the same values as `evaluate`, including collections: + + >>> fvs = fvsof([a(), {'k': f(0, 1)}]) + >>> assert fvs >= {f, a} + """ from effectful.internals.product_n import _unpack, argsof, productN from effectful.internals.runtime import interpreter From 640ec2207c11114599de97d16c66e33a8f8c5796 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Thu, 23 Jul 2026 14:16:58 -0400 Subject: [PATCH 15/26] ensure dataclass constr operations are not returned by fvsof --- effectful/ops/semantics.py | 5 +++-- effectful/ops/types.py | 12 ++---------- tests/test_ops_semantics.py | 2 +- 3 files changed, 6 insertions(+), 13 deletions(-) diff --git a/effectful/ops/semantics.py b/effectful/ops/semantics.py index 52134c771..50066905e 100644 --- a/effectful/ops/semantics.py +++ b/effectful/ops/semantics.py @@ -360,7 +360,7 @@ def _apply_binders(op, *args, **kwargs): # arguments are available through argsof while this node is analyzed. return _BaseTerm(op) - def _apply_collection_fvs(op, *args, **kwargs): + def _apply_passthrough_fvs(op, *args, **kwargs): return frozenset().union( *(x for x in (*args, *kwargs.values()) if isinstance(x, frozenset)) ) @@ -394,7 +394,8 @@ def _apply_fvs(op, *args, **kwargs): { _fvsof_fvs: { apply: _apply_fvs, - CollectionConstrOperation.__apply__: _apply_collection_fvs, + CollectionConstrOperation.__apply__: _apply_passthrough_fvs, + DataclassConstrOperation.__apply__: _apply_passthrough_fvs, }, _fvsof_binders: { CollectionConstrOperation.__apply__: _apply_collection_binders, diff --git a/effectful/ops/types.py b/effectful/ops/types.py index 106c6c646..ec10d37e4 100644 --- a/effectful/ops/types.py +++ b/effectful/ops/types.py @@ -465,19 +465,11 @@ def __set_name__[T](self, owner: type[T], name: str) -> None: def __get__[T](self, instance: T | None, owner: type[T] | None = None): if hasattr(instance, "__dict__") and hasattr(self, "_name_on_instance"): - from effectful.ops.semantics import ( - DataclassConstrOperation, - _as_type, - fvsof, - ) + from effectful.ops.semantics import fvsof - instance_type: type = type(instance) if self._name_on_instance in instance.__dict__: return instance.__dict__[self._name_on_instance] - elif isinstance(instance, Term) or ( - fvsof(instance) - - {_as_type(instance_type, operation_type=DataclassConstrOperation)} - ): + elif isinstance(instance, Term) or (fvsof(instance)): return types.MethodType(self, instance) else: diff --git a/tests/test_ops_semantics.py b/tests/test_ops_semantics.py index fdc6f1cd7..666e51fe9 100644 --- a/tests/test_ops_semantics.py +++ b/tests/test_ops_semantics.py @@ -801,7 +801,7 @@ def __init__(self, x: int): v = Operation.define(int) actual = fvsof(A(v())) - assert actual >= {v} + assert actual == {v} def test_defdata_dataclass_init_effects() -> None: From 81cd5456c1c70d891a2a50c386d5feb232a3f1ad Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Thu, 23 Jul 2026 14:17:45 -0400 Subject: [PATCH 16/26] reduce diff --- effectful/ops/types.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/effectful/ops/types.py b/effectful/ops/types.py index ec10d37e4..65e31de38 100644 --- a/effectful/ops/types.py +++ b/effectful/ops/types.py @@ -469,7 +469,7 @@ def __get__[T](self, instance: T | None, owner: type[T] | None = None): if self._name_on_instance in instance.__dict__: return instance.__dict__[self._name_on_instance] - elif isinstance(instance, Term) or (fvsof(instance)): + elif isinstance(instance, Term) or fvsof(instance): return types.MethodType(self, instance) else: From a64317a896edb1f8524e74814d2f7505cb6ed1e0 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Sat, 25 Jul 2026 10:04:19 -0400 Subject: [PATCH 17/26] wip --- effectful/handlers/jax/_handlers.py | 4 +- effectful/ops/semantics.py | 77 +++++++++++++++-------------- effectful/ops/syntax.py | 30 +++++------ 3 files changed, 53 insertions(+), 58 deletions(-) diff --git a/effectful/handlers/jax/_handlers.py b/effectful/handlers/jax/_handlers.py index bb3ad8b86..1095dcfef 100644 --- a/effectful/handlers/jax/_handlers.py +++ b/effectful/handlers/jax/_handlers.py @@ -14,7 +14,7 @@ from effectful.internals.runtime import interpreter from effectful.ops.semantics import apply, evaluate, fvsof, typeof from effectful.ops.syntax import ( - CollectionConstrOperation, + ConstructorOperation, Scoped, _BaseTerm, _CustomSingleDispatchCallable, @@ -107,7 +107,7 @@ def _getitem(arr, index): _getitem_term: { apply: _retain, jax_getitem: _retain_getitem, - CollectionConstrOperation.__apply__: apply.__default_rule__, + ConstructorOperation.__apply__: apply.__default_rule__, }, } ) diff --git a/effectful/ops/semantics.py b/effectful/ops/semantics.py index 50066905e..0df0c535d 100644 --- a/effectful/ops/semantics.py +++ b/effectful/ops/semantics.py @@ -9,12 +9,10 @@ from typing import Any from effectful.ops.syntax import ( - CollectionConstrOperation, + ConstructorOperation, + DataclassConstructorOperation, _BaseTerm, _CustomSingleDispatchCallable, - as_list, - as_set, - as_tuple, defop, ) from effectful.ops.types import ( @@ -123,18 +121,6 @@ def handler(intp: Interpretation): yield intp -class DataclassConstrOperation(Operation): ... - - -@functools.cache -def _as_type(typ: type, operation_type=CollectionConstrOperation): - @operation_type.define - def _as_typ(*args, **kwargs) -> typ: # type: ignore[valid-type] - return typ(*args, **kwargs) - - return _as_typ - - @_CustomSingleDispatchCallable def evaluate[T]( __dispatch: Callable[[type], Callable[..., Expr[T]]], @@ -178,13 +164,11 @@ def _evaluate_object[T](expr: T, **kwargs) -> T: def _evaluate_dataclass[T](expr: T, **kwargs) -> T: - typ: type = type(expr) - dataclass_op = _as_type(typ, operation_type=DataclassConstrOperation) subst = { field.name: evaluate(getattr(expr, field.name)) for field in dataclasses.fields(expr) # type: ignore[arg-type] } - return typing.cast(T, dataclass_op(**subst)) + return typing.cast(T, DataclassConstructorOperation.define(type(expr))(**subst)) @evaluate.register(Term) @@ -204,19 +188,24 @@ def _evaluate_operation(expr: Operation, **kwargs) -> Operation: @evaluate.register(collections.defaultdict) def _evaluate_defaultdict(expr, **kwargs): - return _as_type(type(expr))( - expr.default_factory, as_list(*(evaluate(item) for item in expr.items())) + return ConstructorOperation.define(type(expr))( + expr.default_factory, + as_tuple(*(evaluate(item) for item in expr.items())), ) @evaluate.register(types.MappingProxyType) def _evaluate_mappingproxytype(expr, **kwargs): - return _as_type(type(expr))(as_list(*(evaluate(item) for item in expr.items()))) + return ConstructorOperation.define(type(expr))( + as_tuple(*(evaluate(item) for item in expr.items())) + ) @evaluate.register(collections.abc.Mapping) def _evaluate_mapping(expr, **kwargs): - return _as_type(type(expr))(as_list(*(evaluate(item) for item in expr.items()))) + return ConstructorOperation.define(type(expr))( + as_tuple(*(evaluate(item) for item in expr.items())) + ) @evaluate.register(tuple) @@ -226,29 +215,35 @@ def _evaluate_tuple(expr, **kwargs): and hasattr(expr, "_fields") and all(hasattr(expr, field) for field in getattr(expr, "_fields")) ): # namedtuple - return _as_type(type(expr))( + return ConstructorOperation.define(type(expr))( **{field: evaluate(getattr(expr, field)) for field in expr._fields} ) else: - return _as_type(type(expr))(as_tuple(*(evaluate(item) for item in expr))) + return ConstructorOperation.define(type(expr))( + as_tuple(*(evaluate(item) for item in expr)) + ) @evaluate.register(collections.abc.Sequence) def _evaluate_sequence(expr, **kwargs): - seq = as_list(*(evaluate(item) for item in expr)) - cast = _as_type(type(expr)) - return cast(seq) + return ConstructorOperation.define(type(expr))( + as_tuple(*(evaluate(item) for item in expr)) + ) @evaluate.register(collections.abc.ItemsView) @evaluate.register(collections.abc.KeysView) def _evaluate_set_view(expr, **kwargs): - return as_set(*(evaluate(item) for item in expr)) + return ConstructorOperation.define(set)( + as_tuple(*(evaluate(item) for item in expr)) + ) @evaluate.register(collections.abc.ValuesView) def _evaluate_list_view(expr, **kwargs): - return as_list(*(evaluate(item) for item in expr)) + return ConstructorOperation.define(list)( + as_tuple(*(evaluate(item) for item in expr)) + ) def _simple_type(tp: type) -> type: @@ -302,7 +297,11 @@ def _apply(op, *args, **kwargs): return Box(op.__type_rule__(*args, **kwargs)) with interpreter( - {apply: _apply, CollectionConstrOperation.__apply__: apply.__default_rule__} + { + apply: _apply, + ConstructorOperation.__apply__: apply.__default_rule__, + DataclassConstructorOperation.__apply__: _apply, + } ): type_or_value = evaluate(term) if isinstance(type_or_value, Box): @@ -343,7 +342,7 @@ def fvsof[S](term: Expr[S]) -> collections.abc.Set[Operation]: _fvsof_fvs = defop(object, name="fvsof_fvs") _fvsof_binders = defop(object, name="fvsof_binders") - def _apply_collection_binders(op, *args, **kwargs): + def _apply_collection_binders(*args, **kwargs): return frozenset().union( *( {x} @@ -394,13 +393,15 @@ def _apply_fvs(op, *args, **kwargs): { _fvsof_fvs: { apply: _apply_fvs, - CollectionConstrOperation.__apply__: _apply_passthrough_fvs, - DataclassConstrOperation.__apply__: _apply_passthrough_fvs, - }, - _fvsof_binders: { - CollectionConstrOperation.__apply__: _apply_collection_binders, - apply: _apply_binders, + ConstructorOperation.__apply__: _apply_passthrough_fvs, }, + _fvsof_binders: ( + {apply: _apply_binders} + | { + ConstructorOperation.define(t): _apply_collection_binders + for t in (dict, list, set, tuple) + } + ), } ) diff --git a/effectful/ops/syntax.py b/effectful/ops/syntax.py index 2f5a53d2d..493708c7f 100644 --- a/effectful/ops/syntax.py +++ b/effectful/ops/syntax.py @@ -518,11 +518,11 @@ def apply_cast(op, *args, **kwargs): { typ: { apply: apply_type, - CollectionConstrOperation.__apply__: apply.__default_rule__, + ConstructorOperation.__apply__: apply.__default_rule__, }, cast: { apply: apply_cast, - CollectionConstrOperation.__apply__: apply.__default_rule__, + ConstructorOperation.__apply__: apply.__default_rule__, }, } ) @@ -1340,24 +1340,18 @@ class _BoolTerm[T: bool](_IntegralTerm[T]): # type: ignore pass -class CollectionConstrOperation(Operation): ... +class ConstructorOperation[**Q, V](Operation[Q, V]): + @classmethod + @functools.cache + def define[T](cls, typ: type[T]) -> "ConstructorOperation[Any, T]": + @Operation.define + def _as_typ(*args, **kwargs) -> typ: # type: ignore[valid-type] + return typ(*args, **kwargs) + + return _as_typ @CollectionConstrOperation.define def as_tuple(*args) -> tuple: return tuple(args) - - -@CollectionConstrOperation.define -def as_list[T](*args: T) -> list[T]: - return list(args) - - -@CollectionConstrOperation.define -def as_set[T](*args: T) -> set[T]: - return set(args) - - -@CollectionConstrOperation.define -def as_dict[K, V](*args: tuple[K, V]) -> dict[K, V]: - return dict(args) +class DataclassConstructorOperation[**Q, V](ConstructorOperation): ... From 6c8e709957af1421173141f7ab9f54ebeb858451 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Mon, 27 Jul 2026 11:32:42 -0400 Subject: [PATCH 18/26] fix bugs --- effectful/ops/semantics.py | 18 +++++++++--------- effectful/ops/syntax.py | 3 +-- 2 files changed, 10 insertions(+), 11 deletions(-) diff --git a/effectful/ops/semantics.py b/effectful/ops/semantics.py index 1938985f9..2ca4dba6c 100644 --- a/effectful/ops/semantics.py +++ b/effectful/ops/semantics.py @@ -176,7 +176,10 @@ def _evaluate_dataclass[T](expr: T, **kwargs) -> T: field.name: evaluate(getattr(expr, field.name)) for field in dataclasses.fields(expr) # type: ignore[arg-type] } - return typing.cast(T, DataclassConstructorOperation.define(type(expr))(**subst)) + return typing.cast( + T, + DataclassConstructorOperation.define(type(expr))(**subst), # type: ignore[arg-type] + ) _EVALUATION_CACHE_ATTR = "__effectful_evaluation_cache__" @@ -399,7 +402,7 @@ def fvsof[S](term: Expr[S]) -> collections.abc.Set[Operation]: _fvsof_fvs = defop(object, name="fvsof_fvs") _fvsof_binders = defop(object, name="fvsof_binders") - def _apply_collection_binders(*args, **kwargs): + def _apply_collection_binders(op, *args, **kwargs): return frozenset().union( *( {x} @@ -452,13 +455,10 @@ def _apply_fvs(op, *args, **kwargs): apply: _apply_fvs, ConstructorOperation.__apply__: _apply_passthrough_fvs, }, - _fvsof_binders: ( - {apply: _apply_binders} - | { - ConstructorOperation.define(t): _apply_collection_binders - for t in (dict, list, set, tuple) - } - ), + _fvsof_binders: { + apply: _apply_binders, + ConstructorOperation.__apply__: _apply_collection_binders, + }, } ) diff --git a/effectful/ops/syntax.py b/effectful/ops/syntax.py index 82f9ccce6..381d22833 100644 --- a/effectful/ops/syntax.py +++ b/effectful/ops/syntax.py @@ -1366,11 +1366,10 @@ class ConstructorOperation[**Q, V](Operation[Q, V]): @classmethod @functools.cache def define[T](cls, typ: type[T]) -> "ConstructorOperation[Any, T]": - @Operation.define def _as_typ(*args, **kwargs) -> typ: # type: ignore[valid-type] return typ(*args, **kwargs) - return _as_typ + return typing.cast(ConstructorOperation[Any, T], super().define(_as_typ)) class DataclassConstructorOperation[**Q, V](ConstructorOperation): ... From d1e7e1322285f9a48b78545d419bcffcee346338 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Mon, 27 Jul 2026 11:33:59 -0400 Subject: [PATCH 19/26] use cached implementation of typeof --- effectful/ops/semantics.py | 19 ++++--------------- 1 file changed, 4 insertions(+), 15 deletions(-) diff --git a/effectful/ops/semantics.py b/effectful/ops/semantics.py index 2ca4dba6c..143bb1834 100644 --- a/effectful/ops/semantics.py +++ b/effectful/ops/semantics.py @@ -350,23 +350,12 @@ def typeof[T](term: Expr[T]) -> type[T]: """ - from effectful.internals.runtime import interpreter from effectful.internals.unification import Box - def _apply(op, *args, **kwargs): - return Box(op.__type_rule__(*args, **kwargs)) - - with interpreter( - { - apply: _apply, - ConstructorOperation.__apply__: apply.__default_rule__, - DataclassConstructorOperation.__apply__: _apply, - } - ): - type_or_value = evaluate(term) - if isinstance(type_or_value, Box): - return _simple_type(type_or_value.value) - return typing.cast(type[T], type(type_or_value)) + type_or_value = _typeof(term) + if isinstance(type_or_value, Box): + return _simple_type(type_or_value.value) + return typing.cast(type[T], type(type_or_value)) def fvsof[S](term: Expr[S]) -> collections.abc.Set[Operation]: From 7082fd891b5dcbe51fe44170b8d568cedfce725c Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Mon, 27 Jul 2026 11:54:08 -0400 Subject: [PATCH 20/26] cache fvsof --- effectful/ops/semantics.py | 105 +++++++++++++++++++----------------- effectful/ops/syntax.py | 24 ++++++--- tests/test_ops_semantics.py | 29 +++++----- 3 files changed, 90 insertions(+), 68 deletions(-) diff --git a/effectful/ops/semantics.py b/effectful/ops/semantics.py index 143bb1834..dc54cd5f0 100644 --- a/effectful/ops/semantics.py +++ b/effectful/ops/semantics.py @@ -12,6 +12,7 @@ from effectful.ops.syntax import ( ConstructorOperation, DataclassConstructorOperation, + ObjectInterpretation, PureInterpretation, _BaseTerm, _CustomSingleDispatchCallable, @@ -312,7 +313,7 @@ def _simple_type(tp: type) -> type: return typing.get_origin(tp) or tp -class _TypeofIntp(PureInterpretation): +class _TypeofIntp(ObjectInterpretation): @implements(apply) def _(self, op, *args, **kwargs): from effectful.internals.unification import Box @@ -320,7 +321,7 @@ def _(self, op, *args, **kwargs): return Box(op.__type_rule__(*args, **kwargs)) -_TYPEOF_INTP = _TypeofIntp() +_TYPEOF_INTP = PureInterpretation(_TypeofIntp()) def _typeof(term: Expr): @@ -358,38 +359,10 @@ def typeof[T](term: Expr[T]) -> type[T]: return typing.cast(type[T], type(type_or_value)) -def fvsof[S](term: Expr[S]) -> collections.abc.Set[Operation]: - """Return the free operations in a term. - - An operation belongs to `fvsof(t)` when it appears free in the term `t`. - This excludes operations like `apply` or collection constructors that are - raised during `evaluate` but do not appear in `t`. It also excludes - operations that are bound by a `Scoped` operation. However, it is not - restricted to the nullary operations in `t`. - - **Example usage**: - - `fvsof` includes all unbound operations in a term: - - >>> a = defop(int) - >>> @defop - ... def f(x: int, y: int) -> int: - ... raise NotHandled - >>> fvs = fvsof(f(a(), 1)) - >>> assert fvs >= {f, a} - - `fvsof` accepts the same values as `evaluate`, including collections: - - >>> fvs = fvsof([a(), {'k': f(0, 1)}]) - >>> assert fvs >= {f, a} - - """ - from effectful.internals.product_n import _unpack, argsof, productN - from effectful.internals.runtime import interpreter - - # Analysis for type computation and term reconstruction - _fvsof_fvs = defop(object, name="fvsof_fvs") - _fvsof_binders = defop(object, name="fvsof_binders") +@functools.cache +def _fvsof_intp() -> tuple[PureInterpretation, Operation]: + """Construct the singleton interpretation used by ``fvsof``.""" + from effectful.internals.product_n import argsof, productN def _apply_collection_binders(op, *args, **kwargs): return frozenset().union( @@ -438,23 +411,59 @@ def _apply_fvs(op, *args, **kwargs): fvs -= binders return fvs - _fvsof_intp = productN( - { - _fvsof_fvs: { - apply: _apply_fvs, - ConstructorOperation.__apply__: _apply_passthrough_fvs, - }, - _fvsof_binders: { - apply: _apply_binders, - ConstructorOperation.__apply__: _apply_collection_binders, - }, - } + _fvsof_fvs = defop(object, name="fvsof_fvs") + _fvsof_binders = defop(object, name="fvsof_binders") + + return ( + PureInterpretation( + productN( + { + _fvsof_fvs: { + apply: _apply_fvs, + ConstructorOperation.__apply__: _apply_passthrough_fvs, + }, + _fvsof_binders: { + apply: _apply_binders, + ConstructorOperation.__apply__: _apply_collection_binders, + }, + } + ) + ), + _fvsof_fvs, ) - with interpreter(_fvsof_intp): - result = evaluate(term) - fvs = _unpack(result, _fvsof_fvs) +def fvsof[S](term: Expr[S]) -> collections.abc.Set[Operation]: + """Return the free operations in a term. + + An operation belongs to `fvsof(t)` when it appears free in the term `t`. + This excludes operations like `apply` or collection constructors that are + raised during `evaluate` but do not appear in `t`. It also excludes + operations that are bound by a `Scoped` operation. However, it is not + restricted to the nullary operations in `t`. + + **Example usage**: + + `fvsof` includes all unbound operations in a term: + + >>> a = defop(int) + >>> @defop + ... def f(x: int, y: int) -> int: + ... raise NotHandled + >>> fvs = fvsof(f(a(), 1)) + >>> assert fvs >= {f, a} + + `fvsof` accepts the same values as `evaluate`, including collections: + + >>> fvs = fvsof([a(), {'k': f(0, 1)}]) + >>> assert fvs >= {f, a} + + """ + from effectful.internals.product_n import _unpack + + intp, prompt = _fvsof_intp() + result = evaluate(term, intp=intp) + fvs = _unpack(result, prompt) if not isinstance(fvs, frozenset): return frozenset() return fvs diff --git a/effectful/ops/syntax.py b/effectful/ops/syntax.py index 381d22833..4330e6423 100644 --- a/effectful/ops/syntax.py +++ b/effectful/ops/syntax.py @@ -11,6 +11,7 @@ from effectful.ops.types import ( Annotation, Expr, + Interpretation, NotHandled, Operation, Term, @@ -952,14 +953,23 @@ def __getitem__(self, item: Operation[..., T]) -> Callable[..., V]: return self.implementations[item].__get__(self, type(self)) -class PureInterpretation[T, V](ObjectInterpretation[T, V]): - def __hash__(self): - return hash(frozenset(self.implementations.items())) +class PureInterpretation[T, V](Mapping[Operation[..., T], Callable[..., V]]): + """Mark an interpretation as pure so its evaluation results can be cached.""" + + def __init__(self, intp: Interpretation[T, V]): + self.intp = intp + + def __iter__(self): + return iter(self.intp) + + def __len__(self): + return len(self.intp) + + def __getitem__(self, item: Operation[..., T]) -> Callable[..., V]: + return self.intp[item] - def __eq__(self, other): - return isinstance(other, PureInterpretation) and frozenset( - self.implementations.items() - ) == frozenset(other.implementations.items()) + __hash__ = object.__hash__ + __eq__ = object.__eq__ class _ImplementedOperation[**P, **Q, T, V]: diff --git a/tests/test_ops_semantics.py b/tests/test_ops_semantics.py index 301f664fc..bd6c5d0de 100644 --- a/tests/test_ops_semantics.py +++ b/tests/test_ops_semantics.py @@ -481,7 +481,7 @@ def node(x: object) -> object: term = node(node(1)) - class Intp(PureInterpretation): + class Intp(ObjectInterpretation): def __init__(self): self.calls = 0 @@ -490,32 +490,34 @@ def _(self, op, *args, **kwargs): self.calls += 1 return (op.__name__, args, kwargs) - intp = Intp() + intp_impl = Intp() + intp = PureInterpretation(intp_impl) expected = ("node", (("node", (1,), {}),), {}) assert interpreter(intp)(evaluate)(term) == expected - assert intp.calls == 2 + assert intp_impl.calls == 2 # The root cache is checked before its children are traversed, including # when evaluation is expressed directly through a handler. with interpreter(intp): assert evaluate(term) == expected - assert intp.calls == 2 + assert intp_impl.calls == 2 # Child results are cached independently and can be reused directly. assert interpreter(intp)(evaluate)(term.args[0]) == expected[1][0] - assert intp.calls == 2 + assert intp_impl.calls == 2 # A composition has a distinct identity even when its added handler is not # used while evaluating this term. combined_intp = coproduct(intp, {plus_1: lambda x: x}) assert interpreter(combined_intp)(evaluate)(term) == expected - assert intp.calls == 4 + assert intp_impl.calls == 4 - # An identical memoized interpretation is in the same cache namespace. - other_intp = Intp() + # A separate pure interpretation has its own cache namespace. + other_intp_impl = Intp() + other_intp = PureInterpretation(other_intp_impl) assert interpreter(other_intp)(evaluate)(term) == expected - assert intp.calls == 4 + assert other_intp_impl.calls == 2 def test_memoized_interpretation_does_not_cache_failures(): @@ -527,7 +529,7 @@ def node() -> object: term = node() - class Intp(PureInterpretation): + class Intp(ObjectInterpretation): def __init__(self): self.calls = 0 @@ -538,14 +540,15 @@ def _(self, op, *args, **kwargs): raise ValueError("failed analysis") return "success" - intp = Intp() + intp_impl = Intp() + intp = PureInterpretation(intp_impl) with pytest.raises(ValueError, match="failed analysis"): interpreter(intp)(evaluate)(term) assert interpreter(intp)(evaluate)(term) == "success" - assert intp.calls == 2 + assert intp_impl.calls == 2 assert interpreter(intp)(evaluate)(term) == "success" - assert intp.calls == 2 + assert intp_impl.calls == 2 @pytest.mark.parametrize( From d9c901f0e6abc5aaf81c96fc5bdc61c4df9d69ff Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Mon, 27 Jul 2026 11:57:01 -0400 Subject: [PATCH 21/26] cache sizesof --- effectful/handlers/jax/_handlers.py | 67 ++++++++++++++++------------- 1 file changed, 38 insertions(+), 29 deletions(-) diff --git a/effectful/handlers/jax/_handlers.py b/effectful/handlers/jax/_handlers.py index 1095dcfef..9c933af43 100644 --- a/effectful/handlers/jax/_handlers.py +++ b/effectful/handlers/jax/_handlers.py @@ -11,10 +11,10 @@ except ImportError: raise ImportError("JAX is required to use effectful.handlers.jax") -from effectful.internals.runtime import interpreter from effectful.ops.semantics import apply, evaluate, fvsof, typeof from effectful.ops.syntax import ( ConstructorOperation, + PureInterpretation, Scoped, _BaseTerm, _CustomSingleDispatchCallable, @@ -42,23 +42,11 @@ def is_eager_array(x): ) -def sizesof(term: Expr) -> Mapping[Operation[[], jax.Array], int]: - """Return the sizes of named dimensions in an array expression. - - Sizes are inferred from the array shape. - - :param value: An array expression. - :return: A mapping from named dimensions to their sizes. - - **Example usage**: - - >>> a, b = defop(jax.Array, name='a'), defop(jax.Array, name='b') - >>> sizes = sizesof(jax_getitem(jnp.ones((2, 3)), [a(), b()])) - >>> assert sizes[a] == 2 and sizes[b] == 3 - """ - from effectful.internals.product_n import _unpack, argsof, productN +@functools.cache +def _sizesof_intp() -> tuple[PureInterpretation, Operation]: + """Construct the singleton interpretation used by ``sizesof``.""" + from effectful.internals.product_n import argsof, productN - # Analysis for type computation and term reconstruction _sizes = defop(object, name="sizes") _getitem_term = defop(object, name="getitem_args") @@ -101,21 +89,42 @@ def _getitem(arr, index): ) return functools.reduce(_merge, itertools.chain(arg_sizes, sizes), {}) - _intp = productN( - { - _sizes: {apply: _apply_sizes, jax_getitem: _getitem}, - _getitem_term: { - apply: _retain, - jax_getitem: _retain_getitem, - ConstructorOperation.__apply__: apply.__default_rule__, - }, - } + return ( + PureInterpretation( + productN( + { + _sizes: {apply: _apply_sizes, jax_getitem: _getitem}, + _getitem_term: { + apply: _retain, + jax_getitem: _retain_getitem, + ConstructorOperation.__apply__: apply.__default_rule__, + }, + } + ) + ), + _sizes, ) - with interpreter(_intp): - result = evaluate(term) - fvs = _unpack(result, _sizes) +def sizesof(term: Expr) -> Mapping[Operation[[], jax.Array], int]: + """Return the sizes of named dimensions in an array expression. + + Sizes are inferred from the array shape. + + :param value: An array expression. + :return: A mapping from named dimensions to their sizes. + + **Example usage**: + + >>> a, b = defop(jax.Array, name='a'), defop(jax.Array, name='b') + >>> sizes = sizesof(jax_getitem(jnp.ones((2, 3)), [a(), b()])) + >>> assert sizes[a] == 2 and sizes[b] == 3 + """ + from effectful.internals.product_n import _unpack + + intp, prompt = _sizesof_intp() + result = evaluate(term, intp=intp) + fvs = _unpack(result, prompt) if not isinstance(fvs, dict): return {} return fvs From 0a38e0e8231485c01de503c692fbca7151fdd5d4 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Mon, 27 Jul 2026 12:01:22 -0400 Subject: [PATCH 22/26] fix test failures --- effectful/handlers/torch.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/effectful/handlers/torch.py b/effectful/handlers/torch.py index c9b8f640b..57468c116 100644 --- a/effectful/handlers/torch.py +++ b/effectful/handlers/torch.py @@ -14,7 +14,13 @@ from effectful.internals.runtime import interpreter from effectful.internals.tensor_utils import _desugar_tensor_index from effectful.ops.semantics import apply, evaluate, fvsof, handler, typeof -from effectful.ops.syntax import Scoped, defdata, defop, syntactic_eq +from effectful.ops.syntax import ( + ConstructorOperation, + Scoped, + defdata, + defop, + syntactic_eq, +) from effectful.ops.types import Expr, NotHandled, Operation, Term # + An element of a tensor index expression. @@ -74,7 +80,13 @@ def _torch_getitem_sizeof( def _apply(op, *args, **kwargs): return defdata(op, *args, **kwargs) - with interpreter({torch_getitem: _torch_getitem_sizeof, apply: _apply}): + with interpreter( + { + torch_getitem: _torch_getitem_sizeof, + apply: _apply, + ConstructorOperation.__apply__: apply.__default_rule__, + } + ): evaluate(value) return sizes From 2618d56e758515d10f324eb1cc2845d65c8d8002 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Mon, 27 Jul 2026 12:03:30 -0400 Subject: [PATCH 23/26] cache torch sizesof --- effectful/handlers/torch.py | 114 +++++++++++++++++++++++++----------- 1 file changed, 79 insertions(+), 35 deletions(-) diff --git a/effectful/handlers/torch.py b/effectful/handlers/torch.py index 57468c116..818490647 100644 --- a/effectful/handlers/torch.py +++ b/effectful/handlers/torch.py @@ -1,4 +1,5 @@ import functools +import itertools import typing from collections.abc import Callable, Mapping, Sequence from types import EllipsisType @@ -11,12 +12,13 @@ import torch.utils._pytree as pytree -from effectful.internals.runtime import interpreter from effectful.internals.tensor_utils import _desugar_tensor_index from effectful.ops.semantics import apply, evaluate, fvsof, handler, typeof from effectful.ops.syntax import ( ConstructorOperation, + PureInterpretation, Scoped, + _BaseTerm, defdata, defop, syntactic_eq, @@ -40,6 +42,76 @@ def _getitem_ellipsis_and_none( return torch.reshape(x, new_shape), new_key +@functools.cache +def _sizesof_intp() -> tuple[PureInterpretation, Operation]: + """Construct the singleton interpretation used by ``sizesof``.""" + from effectful.internals.product_n import argsof, productN + + sizes = defop(object, name="sizes") + getitem_term = defop(object, name="getitem_args") + + def _retain(op, *args, **kwargs): + # Non-getitem subterms are opaque to this analysis. Keeping their + # arguments would retain the entire input term unnecessarily. + return _BaseTerm(op) + + def _retain_getitem(*args, **kwargs): + return defdata(torch_getitem, *args, **kwargs) + + def _merge(s1, s2): + result = s1.copy() + for k, v in s2.items(): + if k in result and result[k] != v: + raise ValueError( + f"Named index {k} used in incompatible dimensions of size {result[k]} and {v}" + ) + result[k] = v + return result + + def _apply_sizes(op, *args, **kwargs): + analyses = (x for x in (*args, *kwargs.values()) if isinstance(x, dict)) + return functools.reduce(_merge, analyses, {}) + + def _getitem(x, key): + # Inspect this getitem's arguments in the term projection without + # forcing that projection to retain the getitem result. + term_args, _ = argsof(getitem_term) + term_x, term_key = term_args + + arg_sizes = (value for value in (x, key) if isinstance(value, dict)) + if not isinstance(term_x, torch.Tensor): + return functools.reduce(_merge, arg_sizes, {}) + + shape, desugared_key = _desugar_tensor_index(term_x.shape, term_key) + index_sizes = ( + {k.op: shape[i]} + for i, k in enumerate(desugared_key) + if isinstance(k, Term) + and not k.args + and not k.kwargs + and issubclass(typeof(k), torch.Tensor) + ) + return functools.reduce( + _merge, itertools.chain(arg_sizes, index_sizes), {} + ) + + return ( + PureInterpretation( + productN( + { + sizes: {apply: _apply_sizes, torch_getitem: _getitem}, + getitem_term: { + apply: _retain, + torch_getitem: _retain_getitem, + ConstructorOperation.__apply__: apply.__default_rule__, + }, + } + ) + ), + sizes, + ) + + def sizesof(value) -> Mapping[Operation[[], torch.Tensor], int]: """Return the sizes of named dimensions in a tensor expression. @@ -54,41 +126,13 @@ def sizesof(value) -> Mapping[Operation[[], torch.Tensor], int]: >>> sizes = sizesof(torch.ones(2, 3)[a(), b()]) >>> assert sizes[a] == 2 and sizes[b] == 3 """ - sizes: dict[Operation[[], torch.Tensor], int] = {} - - def _torch_getitem_sizeof( - x: Expr[torch.Tensor], key: tuple[Expr[IndexElement], ...] - ) -> Expr[torch.Tensor]: - if isinstance(x, torch.Tensor): - shape, key_ = _desugar_tensor_index(x.shape, key) - - for i, k in enumerate(key_): - if ( - isinstance(k, Term) - and len(k.args) == 0 - and len(k.kwargs) == 0 - and issubclass(typeof(k), torch.Tensor) - ): - if k.op in sizes and sizes[k.op] != shape[i]: - raise ValueError( - f"Named index {k.op} used in incompatible dimensions of size {sizes[k.op]} and {shape[i]}" - ) - sizes[k.op] = shape[i] - - return defdata(torch_getitem, x, key) - - def _apply(op, *args, **kwargs): - return defdata(op, *args, **kwargs) - - with interpreter( - { - torch_getitem: _torch_getitem_sizeof, - apply: _apply, - ConstructorOperation.__apply__: apply.__default_rule__, - } - ): - evaluate(value) + from effectful.internals.product_n import _unpack + intp, prompt = _sizesof_intp() + result = evaluate(value, intp=intp) + sizes = _unpack(result, prompt) + if not isinstance(sizes, dict): + return {} return sizes From c05d932f827dbbf937d91daa270921a33a18f0aa Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Mon, 27 Jul 2026 12:08:02 -0400 Subject: [PATCH 24/26] format --- effectful/handlers/torch.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/effectful/handlers/torch.py b/effectful/handlers/torch.py index 818490647..62fa8742e 100644 --- a/effectful/handlers/torch.py +++ b/effectful/handlers/torch.py @@ -91,9 +91,7 @@ def _getitem(x, key): and not k.kwargs and issubclass(typeof(k), torch.Tensor) ) - return functools.reduce( - _merge, itertools.chain(arg_sizes, index_sizes), {} - ) + return functools.reduce(_merge, itertools.chain(arg_sizes, index_sizes), {}) return ( PureInterpretation( From 1beb1f8f401e8e4e742b2f0e38165337fbbdb966 Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Mon, 27 Jul 2026 13:15:12 -0400 Subject: [PATCH 25/26] fix typeof --- effectful/ops/semantics.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/effectful/ops/semantics.py b/effectful/ops/semantics.py index dc54cd5f0..e721c87a1 100644 --- a/effectful/ops/semantics.py +++ b/effectful/ops/semantics.py @@ -320,6 +320,10 @@ def _(self, op, *args, **kwargs): return Box(op.__type_rule__(*args, **kwargs)) + @implements(ConstructorOperation.__apply__) + def _constructor(self, op, *args, **kwargs): + return op.__default_rule__(*args, **kwargs) + _TYPEOF_INTP = PureInterpretation(_TypeofIntp()) From 65fb3f639e9d8ca137e964e731bb38eed4db0bdb Mon Sep 17 00:00:00 2001 From: Jack Feser Date: Mon, 27 Jul 2026 14:12:24 -0400 Subject: [PATCH 26/26] fix pyro tests --- effectful/handlers/pyro.py | 10 +++++++--- effectful/ops/syntax.py | 13 ++++++++++--- 2 files changed, 17 insertions(+), 6 deletions(-) diff --git a/effectful/handlers/pyro.py b/effectful/handlers/pyro.py index 2e7924210..791564afc 100644 --- a/effectful/handlers/pyro.py +++ b/effectful/handlers/pyro.py @@ -28,7 +28,7 @@ ) from effectful.internals.runtime import interpreter from effectful.ops.semantics import apply, evaluate, handler, typeof -from effectful.ops.syntax import defdata, defop +from effectful.ops.syntax import ConstructorOperation, defdata, defop from effectful.ops.types import NotHandled, Operation, Term @@ -368,7 +368,9 @@ def _to_named(a): return a # Convert to a term in a context that does not evaluate distribution constructors. - with handler({apply: defdata}): + with handler( + {apply: defdata, ConstructorOperation.__apply__: apply.__default_rule__} + ): d = typing.cast(TorchDistribution, evaluate(value)) if not (isinstance(d, Term) and typeof(d) is TorchDistribution): @@ -403,7 +405,9 @@ def _to_positional(a, indices): else: return a - with handler({apply: defdata}): + with handler( + {apply: defdata, ConstructorOperation.__apply__: apply.__default_rule__} + ): d = typing.cast(TorchDistribution, evaluate(value)) if not (isinstance(d, Term) and typeof(d) is TorchDistribution): diff --git a/effectful/ops/syntax.py b/effectful/ops/syntax.py index 4330e6423..65b70be54 100644 --- a/effectful/ops/syntax.py +++ b/effectful/ops/syntax.py @@ -1375,9 +1375,16 @@ class _BoolTerm[T: bool](_IntegralTerm[T]): # type: ignore class ConstructorOperation[**Q, V](Operation[Q, V]): @classmethod @functools.cache - def define[T](cls, typ: type[T]) -> "ConstructorOperation[Any, T]": - def _as_typ(*args, **kwargs) -> typ: # type: ignore[valid-type] - return typ(*args, **kwargs) + def define[T]( + cls, constructor: type[T] | Callable[..., T] + ) -> "ConstructorOperation[Any, T]": + if not isinstance(constructor, type): + return typing.cast( + ConstructorOperation[Any, T], super().define(constructor) + ) + + def _as_typ(*args, **kwargs) -> constructor: # type: ignore[valid-type] + return constructor(*args, **kwargs) return typing.cast(ConstructorOperation[Any, T], super().define(_as_typ))