diff --git a/effectful/handlers/jax/_handlers.py b/effectful/handlers/jax/_handlers.py index 3ce8392db..c224e82e0 100644 --- a/effectful/handlers/jax/_handlers.py +++ b/effectful/handlers/jax/_handlers.py @@ -1,8 +1,6 @@ import functools -import itertools import typing -from collections.abc import Callable, Mapping, Sequence -from types import EllipsisType +from collections.abc import Callable, Mapping from typing import Annotated try: @@ -11,21 +9,18 @@ except ImportError: raise ImportError("JAX is required to use effectful.handlers.jax") -from effectful.ops.semantics import apply, evaluate, fvsof, typeof +from effectful.internals.tensor_utils import IndexElement, _BaseSizesofIntp, _sizesof +from effectful.ops.semantics import fvsof, typeof from effectful.ops.syntax import ( - ConstructorOperation, Scoped, - _BaseTerm, _CustomSingleDispatchCallable, defdata, deffn, defop, + implements, syntactic_eq, ) -from effectful.ops.types import Expr, Interpretation, NotHandled, Operation, Term - -# + An element of an array index expression. -IndexElement = None | int | slice | Sequence[int] | EllipsisType | jax.Array +from effectful.ops.types import Expr, NotHandled, Operation, Term def is_eager_array(x): @@ -41,92 +36,6 @@ def is_eager_array(x): ) -@functools.cache -def _sizesof_intp() -> tuple[Interpretation, 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(jax_getitem, *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): - # 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): - 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), {}) - - return ( - productN( - { - _sizes: {apply: _apply_sizes, jax_getitem: _getitem}, - _getitem_term: { - apply: _retain, - jax_getitem: _retain_getitem, - ConstructorOperation.__apply__: apply.__default_rule__, - }, - } - ), - _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 - - def _partial_eval(t: Expr[jax.Array]) -> Expr[jax.Array]: """Partially evaluate a term with respect to its sized free variables.""" @@ -224,7 +133,7 @@ def _jax_op(*args, **kwargs) -> jax.Array: @_register_jax_op -def jax_getitem(x: jax.Array, key: tuple[IndexElement, ...]) -> jax.Array: +def jax_getitem(x: jax.Array, key: tuple[IndexElement[jax.Array], ...]) -> jax.Array: """Operation for indexing an array. Unlike the standard __getitem__ method, this operation correctly handles indexing with terms. @@ -232,6 +141,38 @@ def jax_getitem(x: jax.Array, key: tuple[IndexElement, ...]) -> jax.Array: return x[tuple(key)] +class _SizesofIntp(_BaseSizesofIntp[jax.Array]): + arr_type: typing.ClassVar[type] = jax.Array + + @classmethod + def _names_dim(cls, op: Operation[[], jax.Array]) -> bool: + return True + + @implements(jax_getitem) + def _jax_getitem(self, arr, key): + return self._getitem(arr, key) + + +_SIZESOF_INTP = _SizesofIntp() + + +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 + """ + return _sizesof(term, analysis=_SIZESOF_INTP) + + @defop @_CustomSingleDispatchCallable def bind_dims[T, A, B]( diff --git a/effectful/handlers/jax/_terms.py b/effectful/handlers/jax/_terms.py index 812062931..54b2b83c9 100644 --- a/effectful/handlers/jax/_terms.py +++ b/effectful/handlers/jax/_terms.py @@ -7,14 +7,13 @@ import effectful.handlers.jax.numpy as jnp from effectful.handlers.jax._handlers import ( - IndexElement, _partial_eval, _register_jax_op, bind_dims, jax_getitem, unbind_dims, ) -from effectful.internals.tensor_utils import _desugar_tensor_index +from effectful.internals.tensor_utils import IndexElement, _desugar_tensor_index from effectful.ops.syntax import defdata from effectful.ops.types import Expr, NotHandled, Operation, Term @@ -87,7 +86,8 @@ def kwargs(self) -> dict: return self._kwargs def __getitem__( - self, key: Expr[IndexElement] | tuple[Expr[IndexElement], ...] + self, + key: Expr[IndexElement[jax.Array]] | tuple[Expr[IndexElement[jax.Array]], ...], ) -> Expr[jax.Array]: return jax_getitem(self, key if isinstance(key, tuple) else (key,)) diff --git a/effectful/handlers/jax/numpy/__init__.py b/effectful/handlers/jax/numpy/__init__.py index cc20d7498..28bb0c30f 100644 --- a/effectful/handlers/jax/numpy/__init__.py +++ b/effectful/handlers/jax/numpy/__init__.py @@ -10,26 +10,15 @@ for name, op in jax.numpy.__dict__.items(): if isinstance(op, types.ModuleType): continue - - # copy constants - if isinstance(op, float | types.NoneType): + elif isinstance(op, float | types.NoneType): globals()[name] = op - - if callable(op): + elif 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 # Tell mypy about our wrapped functions. if TYPE_CHECKING: diff --git a/effectful/handlers/torch.py b/effectful/handlers/torch.py index ffd65b31a..a6ac4e1d7 100644 --- a/effectful/handlers/torch.py +++ b/effectful/handlers/torch.py @@ -1,8 +1,6 @@ import functools -import itertools import typing from collections.abc import Callable, Mapping, Sequence -from types import EllipsisType from typing import Annotated, Any try: @@ -12,25 +10,27 @@ import torch.utils._pytree as pytree -from effectful.internals.tensor_utils import _desugar_tensor_index -from effectful.ops.semantics import apply, evaluate, fvsof, handler, typeof +from effectful.internals.tensor_utils import ( + IndexElement, + _BaseSizesofIntp, + _desugar_tensor_index, + _sizesof, +) +from effectful.ops.semantics import evaluate, fvsof, handler, typeof from effectful.ops.syntax import ( - ConstructorOperation, Scoped, _BaseTerm, defdata, defop, + implements, syntactic_eq, ) -from effectful.ops.types import Expr, Interpretation, NotHandled, Operation, Term - -# + An element of a tensor index expression. -IndexElement = None | int | slice | Sequence[int] | EllipsisType | torch.Tensor +from effectful.ops.types import Expr, NotHandled, Operation, Term def _getitem_ellipsis_and_none( - x: torch.Tensor, key: tuple[IndexElement, ...] -) -> tuple[torch.Tensor, tuple[IndexElement, ...]]: + x: torch.Tensor, key: tuple[IndexElement[torch.Tensor], ...] +) -> tuple[torch.Tensor, tuple[IndexElement[torch.Tensor], ...]]: """Eliminate ellipses and None in an index expression x[key]. Returns x1, key1 such that x1[key1] == x[key] nand key1 does not contain None or Ellipsis. @@ -41,96 +41,6 @@ def _getitem_ellipsis_and_none( return torch.reshape(x, new_shape), new_key -@functools.cache -def _sizesof_intp() -> tuple[Interpretation, 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 ( - 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. - - Sizes are inferred from the tensor shape. - - :param value: A tensor expression. - :return: A mapping from named dimensions to their sizes. - - **Example usage**: - - >>> a, b = defop(torch.Tensor, name='a'), defop(torch.Tensor, name='b') - >>> sizes = sizesof(torch.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(value, intp=intp) - sizes = _unpack(result, prompt) - if not isinstance(sizes, dict): - return {} - return sizes - - def _partial_eval(t: Expr[torch.Tensor]) -> Expr[torch.Tensor]: """Partially evaluate a term with respect to its sized free variables.""" @@ -316,7 +226,9 @@ def _torch_op(*args, **kwargs) -> torch.Tensor: @_register_torch_op -def torch_getitem(x: torch.Tensor, key: tuple[IndexElement, ...]) -> torch.Tensor: +def torch_getitem( + x: torch.Tensor, key: tuple[IndexElement[torch.Tensor], ...] +) -> torch.Tensor: """Operation for indexing a tensor. .. note:: @@ -369,14 +281,48 @@ def torch_getitem(x: torch.Tensor, key: tuple[IndexElement, ...]) -> torch.Tenso return torch.ops.aten.index(x, tuple(key_l)) +class _SizesofIntp(_BaseSizesofIntp[torch.Tensor]): + arr_type: typing.ClassVar[type] = torch.Tensor + + # Restates the one condition ``_embed_tensor`` imposes that its jax + # counterpart does not. It has to keep agreeing with it. + + @classmethod + def _names_dim(cls, op: Operation) -> bool: + return issubclass(typeof(_BaseTerm(op)), cls.arr_type) + + @implements(torch_getitem) + def _torch_getitem(self, x, key): + return self._getitem(x, key) + + +_SIZESOF_INTP = _SizesofIntp() + + +def sizesof(value) -> Mapping[Operation[[], torch.Tensor], int]: + """Return the sizes of named dimensions in a tensor expression. + + Sizes are inferred from the tensor shape. + + :param value: A tensor expression. + :return: A mapping from named dimensions to their sizes. + + **Example usage**: + + >>> a, b = defop(torch.Tensor, name='a'), defop(torch.Tensor, name='b') + >>> sizes = sizesof(torch.ones(2, 3)[a(), b()]) + >>> assert sizes[a] == 2 and sizes[b] == 3 + """ + return _sizesof(value, analysis=_SIZESOF_INTP) + + @defdata.register(torch.Tensor) def _embed_tensor(ty, op, *args, **kwargs): if ( op is torch_getitem and not isinstance(args[0], Term) - and len(args[1]) > 0 and all( - typeof(k) is torch.Tensor and not k.args and not k.kwargs + issubclass(typeof(k), torch.Tensor) and not k.args and not k.kwargs for k in args[1] if isinstance(k, Term) ) @@ -422,7 +368,9 @@ def kwargs(self) -> dict: return self._kwargs def __getitem__( - self, key: Expr[IndexElement] | tuple[Expr[IndexElement], ...] + self, + key: Expr[IndexElement[torch.Tensor]] + | tuple[Expr[IndexElement[torch.Tensor]], ...], ) -> Expr[torch.Tensor]: return torch_getitem(self, key if isinstance(key, tuple) else (key,)) @@ -546,17 +494,19 @@ def __iter__(self): @Term.register class _EagerTensorTerm(torch.Tensor): - args: tuple[torch.Tensor, tuple[IndexElement, ...]] + args: tuple[torch.Tensor, tuple[IndexElement[torch.Tensor], ...]] kwargs: Mapping[str, object] = {} __match_args__ = ("op", "args", "kwargs") - def __new__(cls, x: torch.Tensor, key: tuple[IndexElement, ...]): + def __new__(cls, x: torch.Tensor, key: tuple[IndexElement[torch.Tensor], ...]): assert not isinstance(x, Term) for k in key: if isinstance(k, Term): - assert typeof(k) is torch.Tensor and not k.args and not k.kwargs + assert ( + issubclass(typeof(k), torch.Tensor) and not k.args and not k.kwargs + ) x, key = _getitem_ellipsis_and_none(x, key) ret = x.as_subclass(cls) diff --git a/effectful/internals/product_n.py b/effectful/internals/product_n.py deleted file mode 100644 index ca7c2f0b3..000000000 --- a/effectful/internals/product_n.py +++ /dev/null @@ -1,224 +0,0 @@ -import collections.abc -import dataclasses -import functools -import types -from collections.abc import Callable, Mapping -from typing import Any - -from effectful.ops.semantics import apply, coproduct, handler -from effectful.ops.syntax import defop -from effectful.ops.types import ( - Interpretation, - NotHandled, # noqa: F401 - Operation, -) - - -@dataclasses.dataclass -class CallByNeed[**P, T]: - func: Callable[P, T] - args: Any # P.args - kwargs: Any # P.kwargs - value: T | None = None - initialized: bool = False - - def __init__(self, func, *args, **kwargs): - self.func = func - self.args = args - self.kwargs = kwargs - - def __call__(self): - if not self.initialized: - self.value = self.func(*self.args, **self.kwargs) - self.initialized = True - return self.value - - -@defop -def argsof(op: Operation) -> tuple[list, dict]: - raise RuntimeError("Prompt argsof not bound.") - - -class Product: - values: object - - def __init__(self, values): - self.values = values - - -def _pack(intp): - from effectful.internals.runtime import interpreter - - return Product(interpreter(intp)(lambda x: x())) - - -def _unpack(x, prompt): - if isinstance(x, Product): - return x.values(prompt) - return x - - -def map_structure(func, expr, _cache=None): - if _cache is None: - _cache = {} - - key = id(expr) - if key in _cache: - ref, result = _cache[key] - if ref is expr: - return result - - def recurse(x): - return map_structure(func, x, _cache) - - if isinstance(expr, collections.abc.Mapping): - if isinstance(expr, collections.defaultdict): - result = type(expr)(expr.default_factory, recurse(tuple(expr.items()))) - elif isinstance(expr, types.MappingProxyType): - result = type(expr)(dict(recurse(tuple(expr.items())))) - else: - result = type(expr)(recurse(tuple(expr.items()))) - elif isinstance(expr, collections.abc.Sequence): - if isinstance(expr, str | bytes): - result = expr - elif ( - isinstance(expr, tuple) - and hasattr(expr, "_fields") - and all(hasattr(expr, field) for field in getattr(expr, "_fields")) - ): # namedtuple - result = type(expr)( - **{field: recurse(getattr(expr, field)) for field in expr._fields} - ) - else: - result = type(expr)(recurse(item) for item in expr) - elif isinstance(expr, collections.abc.Set): - if isinstance(expr, collections.abc.ItemsView | collections.abc.KeysView): - result = {recurse(item) for item in expr} - else: - result = type(expr)(recurse(item) for item in expr) - elif isinstance(expr, collections.abc.ValuesView): - result = [recurse(item) for item in expr] - elif dataclasses.is_dataclass(expr) and not isinstance(expr, type): - result = dataclasses.replace( - expr, - **{ - field.name: recurse(getattr(expr, field.name)) - for field in dataclasses.fields(expr) - }, - ) - else: - result = func(expr) - - _cache[key] = (expr, result) - return result - - -def productN(intps: Mapping[Operation, Interpretation]) -> Interpretation: - # The resulting interpretation supports ops that exist in at least one input - # interpretation - result_ops = set(op for intp in intps.values() for op in intp) - if result_ops is None: - return {} - - renaming = {(prompt, op): defop(op) for prompt in intps for op in result_ops} - - # We enforce isolation between the named interpretations by giving every - # operation a fresh name and giving each operation a translation from - # the fresh names back to the names from their interpretation. - # - # E.g. { a: { f, g }, b: { f, h } } => - # { handler({f: f_a, g: g_a, h: h_default})(f_a), handler({f: f_a, g: g_a})(g_a), - # handler({f: f_b, h: h_b})(f_b), handler({f: f_b, h: h_b})(h_b) } - translation_intps: dict[Operation, Interpretation] = { - prompt: {op: renaming[(prompt, op)] for op in result_ops} for prompt in intps - } - - # For every prompt, build an isolated interpretation that binds all operations. - isolated_intps = { - prompt: { - renaming[(prompt, op)]: handler(translation_intps[prompt])(func) - for op, func in intp.items() - } - for prompt, intp in intps.items() - } - - def product_op(op, *args, **kwargs): - """Compute the product of operation `op` in named interpretations - `intps`. The product operation consumes product arguments and - returns product results. These products are represented as - interpretations. - - """ - assert isinstance(op, Operation) - - result_intp = {} - - def argsof_direct_call(prompt): - return result_intp[prompt].args, result_intp[prompt].kwargs - - def argsof_apply(prompt): - return result_intp[prompt].args[2:], result_intp[prompt].kwargs - - # Every prompt gets an argsof implementation. The implementation is - # either for a direct call to a handler or for a call to an apply - # handler. - argsof_prompts = {} - - for prompt, intp in intps.items(): - # Args and kwargs are expected to be either interpretations with - # bindings for each named analysis in intps or concrete values. - # `get_for_intp` extracts the value that corresponds to this - # analysis. - # - # TODO: `get_for_intp` has to guess whether a dict value is an - # interpretation or not. This is probably a latent bug. - intp_args, intp_kwargs = map_structure( - lambda x: _unpack(x, prompt), (args, kwargs) - ) - - # Making result a CallByNeed has two functions. It avoids some - # work when the result is not requested and it delays evaluation - # so that when the result is requested in `get_for_intp`, it - # evaluates in a context that binds the results of the other - # named interpretations. - isolated_intp = isolated_intps[prompt] - renamed_op = renaming[(prompt, op)] - if op in intp: - result = CallByNeed( - handler(isolated_intp)(renamed_op), *intp_args, **intp_kwargs - ) - argsof_impl = argsof_direct_call - elif apply in intp: - result = CallByNeed( - handler(isolated_intp)(renaming[(prompt, apply)]), - renamed_op, - *intp_args, - **intp_kwargs, - ) - argsof_impl = argsof_apply - else: - # TODO: If an intp does not handle an operation and has no apply - # handler, use the default rule. In the future, we would like to - # instead defer to the enclosing interpretation. This is - # difficult right now, because the output interpretation handles - # all operations with product handlers which would have to be - # skipped over. - result = CallByNeed( - handler(coproduct(isolated_intp, translation_intps[prompt]))( - op.__default_rule__ - ), - *intp_args, - **intp_kwargs, - ) - argsof_impl = argsof_direct_call - - result_intp[prompt] = result - argsof_prompts[prompt] = argsof_impl - - result_intp[argsof] = lambda prompt: argsof_prompts[prompt](prompt) - return _pack(result_intp) - - product_intp: Interpretation = { - op: functools.partial(product_op, op) for op in result_ops - } - return product_intp diff --git a/effectful/internals/tensor_utils.py b/effectful/internals/tensor_utils.py index cbf187915..6d3d692e4 100644 --- a/effectful/internals/tensor_utils.py +++ b/effectful/internals/tensor_utils.py @@ -1,8 +1,25 @@ -def _desugar_tensor_index(shape, key): - new_shape = [] - new_key = [] +import abc +import collections.abc +import functools +import types +import typing - def extra_dims(key): +from effectful.ops.semantics import ConstructorOperation, apply, evaluate +from effectful.ops.syntax import ObjectInterpretation, implements +from effectful.ops.types import Operation + +type IndexElement[T] = ( + None | int | slice | collections.abc.Sequence[int] | types.EllipsisType | T +) + + +def _desugar_tensor_index[T]( + shape: tuple[int, ...], key: collections.abc.Sequence[IndexElement[T]] +) -> tuple[tuple[int, ...], tuple[IndexElement[T], ...]]: + new_shape: list[int] = [] + new_key: list[IndexElement[T]] = [] + + def extra_dims(key: collections.abc.Sequence[IndexElement[T]]) -> int: return sum(1 for k in key if k is None) # handle any missing dimensions by adding a trailing Ellipsis @@ -29,4 +46,139 @@ def extra_dims(key): new_shape.append(shape[len(new_shape) - extra_dims(key[:i])]) new_key.append(k) - return new_shape, new_key + return tuple(new_shape), tuple(new_key) + + +class _Name[T]: + """An index entry that names a dimension: a bare call to ``op``. + + Deliberately not a tuple, so that a key can be told apart from an entry. + """ + + __slots__ = ("op",) + + def __init__(self, op: Operation[[], T]): + self.op = op + + +#: An index entry that is a term but not a bare name, so it neither names a +#: dimension nor leaves the indexed result with a shape this analysis can +#: predict. Distinct from a concrete entry, which does neither but is harmless. +_OPAQUE: typing.Any = object() + + +class _SizeAnalysis[T](typing.NamedTuple): + """What the analysis of a single node carries. + + ``sizes`` is the result. The rest is what a parent `__getitem__` + needs to finish its own analysis, which the sizes alone cannot supply: + ``shape`` when the node denotes an array whose shape is known, and + ``index`` for what the node looks like in a key -- the dimension it names, + or the value it already is. + """ + + sizes: dict[Operation[[], T], int] + index: typing.Any + shape: tuple[int, ...] | None = None + + +class _BaseSizesofIntp[T](abc.ABC, ObjectInterpretation): + """Shared part of the analysis behind ``sizesof``. + + The hook below is where the backends differ, and it has to answer exactly + as that backend's :func:`defdata` rule does. That rule decides whether an + indexed result is built eagerly, and so whether it has a shape at all; an + analysis that disagreed would predict a shape for a term that is never + built with one, or miss one that is. The default is the permissive answer, + which is what ``_embed_array`` gives; ``_embed_tensor`` is stricter about + what may name a dimension and overrides it. + """ + + arr_type: typing.ClassVar[type] = object + + @classmethod + @abc.abstractmethod + def _names_dim(cls, op: Operation[[], T]) -> bool: + """Whether a bare call to ``op`` names a dimension of what it indexes.""" + raise NotImplementedError + + @classmethod + def _analysis(cls, value) -> _SizeAnalysis[T]: + """View a rule argument as an analysis. Leaves contribute no sizes. + + A leaf stands for itself in a key, so that keys rebuild into real + tuples holding real slices and ``None`` and ``Ellipsis`` literals. + """ + if isinstance(value, _SizeAnalysis): + return value + elif isinstance(value, cls.arr_type): + return _SizeAnalysis[T]({}, value, value.shape) # type: ignore + else: + return _SizeAnalysis[T]({}, value) + + @staticmethod + def _merge( + s1: dict[Operation[[], T], int], s2: dict[Operation[[], T], int] + ) -> dict[Operation[[], T], int]: + 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 + + @implements(apply) + def _apply(self, op, *args, **kwargs): + analyses = tuple(self._analysis(x) for x in (*args, *kwargs.values())) + return _SizeAnalysis( + functools.reduce(self._merge, (a.sizes for a in analyses), {}), + _Name(op) if not (args or kwargs) and self._names_dim(op) else _OPAQUE, + ) + + @implements(ConstructorOperation.__apply__) + def _apply_constructor(self, op, *args, **kwargs): + arg_analyses = tuple(self._analysis(x) for x in args) + kwarg_analyses = {k: self._analysis(v) for k, v in kwargs.items()} + analyses = (*arg_analyses, *kwarg_analyses.values()) + return _SizeAnalysis( + functools.reduce(self._merge, (a.sizes for a in analyses), {}), + op.__default_rule__( + *(a.index for a in arg_analyses), + **{k: a.index for k, a in kwarg_analyses.items()}, + ), + ) + + def _getitem(self, x, key): + is_concrete = isinstance(x, self.arr_type) + x, key = self._analysis(x), self._analysis(key) + sizes = self._merge(x.sizes, key.sizes) + + if x.shape is None or not isinstance(key.index, tuple | list): + return _SizeAnalysis(sizes, _OPAQUE) + + shape, entries = _desugar_tensor_index(x.shape, key.index) + for i, entry in enumerate(entries): + if isinstance(entry, _Name): + sizes = self._merge(sizes, {entry.op: shape[i]}) + + eager = is_concrete and not any(e is _OPAQUE for e in entries) + return _SizeAnalysis( + sizes, + _OPAQUE, + tuple(s for s, e in zip(shape, entries) if not isinstance(e, _Name)) + if eager + else None, + ) + + +def _sizesof[T]( + value, *, analysis: _BaseSizesofIntp[T] +) -> collections.abc.Mapping[Operation[[], T], int]: + """Return a mapping from named dimensions to their sizes. + + Raises a ValueError if the same name is used for different sizes. + """ + result = evaluate(value, intp=analysis) + return result.sizes if isinstance(result, _SizeAnalysis) else {} diff --git a/effectful/ops/semantics.py b/effectful/ops/semantics.py index 6fac67e35..33f386c58 100644 --- a/effectful/ops/semantics.py +++ b/effectful/ops/semantics.py @@ -11,7 +11,7 @@ ConstructorOperation, DataclassConstructorOperation, ObjectInterpretation, - _BaseTerm, + Scoped, _CustomSingleDispatchCallable, defop, implements, @@ -322,11 +322,6 @@ def _dataclass_constructor_apply(self, op, *args, **kwargs): _TYPEOF_INTP = _TypeofIntp() -def _typeof(term: Expr): - """Evaluate the cached type analysis without unwrapping its result.""" - return evaluate(term, intp=_TYPEOF_INTP) - - def typeof[T](term: Expr[T]) -> type[T]: """Return the type of an expression. @@ -351,82 +346,49 @@ def typeof[T](term: Expr[T]) -> type[T]: """ from effectful.internals.unification import Box - type_or_value = _typeof(term) + type_or_value = evaluate(term, intp=_TYPEOF_INTP) if isinstance(type_or_value, Box): return _simple_type(type_or_value.value) return typing.cast(type[T], type(type_or_value)) -@functools.cache -def _fvsof_intp() -> tuple[Interpretation, 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( - *( - {x} - if isinstance(x, Operation) - else x - if isinstance(x, frozenset) - else set() - for x in (*args, *kwargs.values()) - ) - ) +class _FvsAnalysis(typing.NamedTuple): + ops: frozenset[Operation] = frozenset() + fvs: frozenset[Operation] = frozenset() + - def _apply_binders(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) +class _FvsofIntp(ObjectInterpretation): + @staticmethod + def _analysis(value) -> _FvsAnalysis: + if isinstance(value, _FvsAnalysis): + return value + else: + return _FvsAnalysis(Scoped.extract_operations(value)) - def _apply_passthrough_fvs(op, *args, **kwargs): - return frozenset().union( - *(x for x in (*args, *kwargs.values()) if isinstance(x, frozenset)) + @implements(ConstructorOperation.__apply__) + def _apply_collection_binders(self, op, *args, **kwargs): + analyses = tuple(self._analysis(x) for x in (*args, *kwargs.values())) + return _FvsAnalysis( + frozenset().union(frozenset(), *(a.ops for a in analyses)), + frozenset().union(frozenset(), *(a.fvs for a in analyses)), ) - def _apply_fvs(op, *args, **kwargs): - 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:] + @implements(apply) + def _apply_fvs(self, op, *args, **kwargs): + arg_analyses = tuple(self._analysis(a) for a in args) + kwarg_analyses = {k: self._analysis(v) for k, v in kwargs.items()} + bindings = op.__fvs_rule__( + *(a.ops for a in arg_analyses), + **{k: a.ops for k, a in kwarg_analyses.items()}, ) - 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}, - *( - x if isinstance(x, frozenset) else frozenset() - for x in (*args, *kwargs.values()) - ), + {op}, *(a.fvs for a in (*arg_analyses, *kwarg_analyses.values())) ) - fvs -= binders - return fvs - - _fvsof_fvs = defop(object, name="fvsof_fvs") - _fvsof_binders = defop(object, name="fvsof_binders") - - return ( - productN( - { - _fvsof_fvs: { - apply: _apply_fvs, - ConstructorOperation.__apply__: _apply_passthrough_fvs, - }, - _fvsof_binders: { - apply: _apply_binders, - ConstructorOperation.__apply__: _apply_collection_binders, - }, - } - ), - _fvsof_fvs, - ) + return _FvsAnalysis(fvs=fvs - binders) + + +_FVSOF_INTP = _FvsofIntp() def fvsof[S](term: Expr[S]) -> collections.abc.Set[Operation]: @@ -455,11 +417,5 @@ def fvsof[S](term: Expr[S]) -> collections.abc.Set[Operation]: >>> 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 + result = evaluate(term, intp=_FVSOF_INTP) + return result.fvs if isinstance(result, _FvsAnalysis) else frozenset() diff --git a/effectful/ops/syntax.py b/effectful/ops/syntax.py index 5b372f0f1..c80ac2103 100644 --- a/effectful/ops/syntax.py +++ b/effectful/ops/syntax.py @@ -317,6 +317,41 @@ def infer_annotations(cls, sig: inspect.Signature) -> inspect.Signature: assert cls._get_root_ordinal(inferred_sig) == root_ordinal != set() return inferred_sig + @classmethod + def extract_operations( + cls, value, _seen: set[int] | None = None + ) -> frozenset[Operation]: + """Computes the set of :class:`Operation` s appearing directly in ``value`` . + + An :class:`Operation` counts when it appears as a value in a collection, + including as the key of a mapping, which is why this cannot be written + in terms of :func:`flatten` . It does not count when it is applied to + arguments, since the resulting :class:`Term` is a use of the operation + rather than a binding occurrence of it. + + :param value: The value to traverse. + :returns: The operations that could be bound by a parameter given ``value``. + """ + _seen = set() if _seen is None else _seen + if id(value) in _seen: + return frozenset() + _seen.add(id(value)) + + if isinstance(value, Operation): + return frozenset({value}) + elif isinstance(value, dict): + return frozenset().union( + frozenset(), + *(cls.extract_operations(k, _seen) for k in value.keys()), + *(cls.extract_operations(v, _seen) for v in value.values()), + ) + elif isinstance(value, list | set | frozenset | tuple): + return frozenset().union( + frozenset(), *(cls.extract_operations(v, _seen) for v in value) + ) + else: + return frozenset() + def analyze(self, bound_sig: inspect.BoundArguments) -> frozenset[Operation]: """ Computes a set of bound variables given a signature with bound arguments. @@ -342,43 +377,19 @@ def analyze(self, bound_sig: inspect.BoundArguments) -> frozenset[Operation]: param_ordinal = self._get_param_ordinal(param) if param_ordinal <= self.ordinal and not param_ordinal <= return_ordinal: param_value = bound_sig.arguments[name] - param_bound_vars = set() - - if self._param_is_var(param): - # Handle individual Operation parameters (existing behavior) - if param.kind is inspect.Parameter.VAR_POSITIONAL: - # pre-condition: all bound variables should be distinct - assert len(param_value) == len(set(param_value)) - param_bound_vars = set(param_value) - elif param.kind is inspect.Parameter.VAR_KEYWORD: - # pre-condition: all bound variables should be distinct - assert len(param_value.values()) == len( - set(param_value.values()) - ) - param_bound_vars = set(param_value.values()) - else: - param_bound_vars = {param_value} - elif param_ordinal: # Only process if there's a Scoped annotation - # We can't use flatten here because we want to be able - # to see dict keys - def extract_operations(obj, _seen=None): - if _seen is None: - _seen = set() - obj_id = id(obj) - if obj_id in _seen: - return - _seen.add(obj_id) - if isinstance(obj, Operation): - param_bound_vars.add(obj) - elif isinstance(obj, dict): - for k, v in obj.items(): - extract_operations(k, _seen) - extract_operations(v, _seen) - elif isinstance(obj, list | set | tuple): - for v in obj: - extract_operations(v, _seen) - - extract_operations(param_value) + param_bound_vars: frozenset[Operation] = ( + self.extract_operations(param_value) + # only process if the parameter is an Operation or is Scoped + if self._param_is_var(param) or param_ordinal + else frozenset() + ) + + if self._param_is_var(param) and param.kind in ( + inspect.Parameter.VAR_POSITIONAL, + inspect.Parameter.VAR_KEYWORD, + ): + # pre-condition: all bound variables should be distinct + assert len(param_bound_vars) == len(param_value) # pre-condition: all bound variables should be distinct if param_bound_vars: diff --git a/tests/test_handlers_jax.py b/tests/test_handlers_jax.py index e07385725..5908f0594 100644 --- a/tests/test_handlers_jax.py +++ b/tests/test_handlers_jax.py @@ -253,6 +253,68 @@ def test_jax_nested_getitem(): assert sizesof(t_ij) == {i: 2, j: 3} +def test_sizesof_compound_index(): + """A compound index is not a named dimension. + + Only a bare call names a dimension. An index built out of one, like + ``a() + 1``, names nothing -- neither the operation applied nor the + variable underneath it. The entry still holds its place, so the names + beside it are found against the right dimensions. + """ + a, b = defop(jax.Array, name="a"), defop(jax.Array, name="b") + + assert sizesof(jax_getitem(jnp.ones((4, 5)), (a() + 1, b()))) == {b: 5} + assert sizesof(jax_getitem(jnp.ones((4, 5)), (a() + 1, a() * 2))) == {} + assert sizesof(jax_getitem(jnp.ones((2, 3)), (a() + 1, None, b()))) == {b: 3} + assert sizesof(jax_getitem(jnp.ones((2, 3)), (None, a() + 1, b()))) == {b: 3} + assert sizesof(jax_getitem(jnp.ones((2, 3, 4)), (a() + 1, ..., b()))) == {b: 4} + + # A bare named dimension is still reported, alongside a compound sibling. + assert sizesof(jax_getitem(jnp.ones((4, 5)), (a(), b()))) == {a: 4, b: 5} + + +def test_sizesof_nested_getitem(): + """Sizes are found through an inner getitem, using its residual shape. + + ``sizesof`` computes that shape itself rather than reading it off a + rebuilt term, so these pin it against the conditions ``_embed_array`` + builds an eager term under: a concrete array whose every term entry is a + bare name. + """ + a, b = defop(jax.Array, name="a"), defop(jax.Array, name="b") + + # A named dimension is consumed by the indexing; the rest stays. + inner = jax_getitem(jnp.ones((2, 3, 4)), (a(), slice(None), slice(None))) + assert sizesof(jax_getitem(inner, (b(), slice(None)))) == {a: 2, b: 3} + + inner = jax_getitem(jnp.ones((2, 3)), (slice(None), a())) + assert sizesof(jax_getitem(inner, (b(),))) == {a: 3, b: 2} + + # A concrete index keeps its dimension. + inner = jax_getitem(jnp.ones((4, 5)), (jnp.arange(4), a())) + assert sizesof(jax_getitem(inner, (b(),))) == {a: 5, b: 4} + + # Unlike torch, an index is not required to be array-typed. + n = defop(int, name="n") + assert sizesof(jax_getitem(jnp.ones((4, 5)), (n(), a()))) == {n: 4, a: 5} + + +def test_sizesof_symbolic_key(): + """A key can be a term rather than a sequence of index entries. + + Indexing the result of an earlier indexing this way is what reaches the + analysis: the inner term has a shape to read while being a term, which is + what stops ``_embed_array`` from inspecting the key at all. + """ + a = defop(jax.Array, name="a") + key = defop(tuple, name="key") + + inner = jax_getitem(jnp.ones((2, 3, 4)), (a(), slice(None), slice(None))) + assert tuple(inner.shape) == (3, 4) + + assert sizesof(jax_getitem(inner, key())) == {a: 2} + + def test_jax_at_updates(): """Test .at array update functionality for indexed arrays.""" i, j, k = defop(jax.Array), defop(jax.Array), defop(jax.Array) diff --git a/tests/test_handlers_torch.py b/tests/test_handlers_torch.py index 6684e808f..069dfb817 100644 --- a/tests/test_handlers_torch.py +++ b/tests/test_handlers_torch.py @@ -226,6 +226,76 @@ def test_tpe_stack(): ] +def test_sizesof_compound_index(): + """A compound index is not a named dimension. + + Only a bare call names a dimension. An index built out of one, like + ``a() + 1``, names nothing -- neither the operation applied nor the + variable underneath it. + """ + a, b = defop(torch.Tensor, name="a"), defop(torch.Tensor, name="b") + + assert sizesof(torch_getitem(torch.ones(4, 5), (a() + 1, b()))) == {b: 5} + assert sizesof(torch_getitem(torch.ones(4, 5), (a() + 1, a() * 2))) == {} + assert sizesof(torch_getitem(torch.ones(4, 5), ((a() + 1) * 2, b()))) == {b: 5} + + # A bare named dimension is still reported, alongside a compound sibling. + assert sizesof(torch.ones(4, 5)[a(), b()]) == {a: 4, b: 5} + + +def test_sizesof_nested_getitem(): + """Sizes are found through an inner getitem, using its residual shape. + + ``sizesof`` computes that shape itself rather than reading it off a rebuilt + term, so these pin it against the conditions ``_embed_tensor`` builds an + eager term under: a concrete tensor indexed by a non-empty key whose every + term entry is a bare, tensor-typed name. + """ + a, b = defop(torch.Tensor, name="a"), defop(torch.Tensor, name="b") + + # A named dimension is consumed by the indexing; the rest stays. + assert sizesof(torch.ones(2, 3, 4)[a(), :, :][b(), :]) == {a: 2, b: 3} + assert sizesof(torch.ones(2, 3)[:, a()][b()]) == {a: 3, b: 2} + + # A concrete index keeps its dimension. + inner = torch_getitem(torch.ones(4, 5), (torch.arange(4), a())) + assert sizesof(torch_getitem(inner, (b(),))) == {a: 5, b: 4} + + # A subclass of torch.Tensor names a dimension just as torch.Tensor does, + # down to consuming it: with every dimension named there is nothing left to + # index into, and both fail alike. + class _SubTensor(torch.Tensor): + pass + + sub = defop(_SubTensor, name="sub") + assert sizesof(torch_getitem(torch.ones(4, 5), (sub(), a()))) == {sub: 4, a: 5} + for name in (sub, defop(torch.Tensor, name="t")): + with pytest.raises(IndexError): + torch_getitem(torch_getitem(torch.ones(4, 5), (name(), a())), (b(),)) + + # An index that is not tensor-typed at all names nothing. + n = defop(int, name="n") + assert sizesof(torch_getitem(torch.ones(4, 5), (n(), a()))) == {a: 5} + + +def test_sizesof_symbolic_key(): + """A key can be a term rather than a sequence of index entries. + + Indexing a concrete tensor this way fails outright, but indexing the + *result* of an earlier indexing does not: the inner term has a shape to + read while being a term, which is what stops ``_embed_tensor`` from + inspecting the key at all. ``sizesof`` still reports what the inner + indexing named. + """ + a = defop(torch.Tensor, name="a") + key = defop(tuple, name="key") + + inner = torch.ones(2, 3, 4)[a(), :, :] + assert inner.shape == (3, 4) + + assert sizesof(torch_getitem(inner, key())) == {a: 2} + + @pytest.mark.parametrize("tensor, idx", INDEXING_CASES) def test_getitem_ellipsis_and_none(tensor, idx): from effectful.handlers.torch import _getitem_ellipsis_and_none diff --git a/tests/test_internals_product_n.py b/tests/test_internals_product_n.py deleted file mode 100644 index 331019824..000000000 --- a/tests/test_internals_product_n.py +++ /dev/null @@ -1,160 +0,0 @@ -from effectful.internals.product_n import argsof, productN -from effectful.internals.unification import Box -from effectful.ops.semantics import apply, coproduct, evaluate, handler -from effectful.ops.syntax import defop -from effectful.ops.types import Interpretation, NotHandled - - -def test_simul_analysis(): - @defop - def plus1(x: int) -> int: - raise NotHandled - - @defop - def plus2(x: int) -> int: - raise NotHandled - - @defop - def times(x: int, y: int) -> int: - raise NotHandled - - x, y = defop(int, name="x"), defop(int, name="y") - - typ = defop(Interpretation, name="typ") - value = defop(Interpretation, name="value") - - type_rules = { - plus1: lambda x: int, - plus2: lambda x: int, - times: lambda x, y: int, - x: lambda: int, - y: lambda: int, - } - - def plus1_value(x): - return x + 1 - - def plus2_value(x): - return plus1(plus1(x)) - - def times_value(x, y): - t = typ() - arg = argsof(typ)[0][0] - if t is int and arg is int: - return x * y - raise TypeError("unexpected type!") - - value_rules = { - plus1: plus1_value, - plus2: plus2_value, - times: times_value, - x: lambda: 3, - y: lambda: 4, - } - - analysisN = productN({typ: type_rules, value: value_rules}) - - def f1(): - v1 = x() # {typ: lambda: int, val: lambda: 3} - v2 = y() # {typ: lambda: int, val: lambda: 4} - v3 = plus2(v1) # {typ: lambda: int, val: lambda: 5} - v4 = times(v2, v3) # {typ: lambda: int, val: lambda: 20} - v5 = plus1(v4) # {typ: lambda: int, val: lambda: 21} - return v5 # {typ: lambda: int, val: lambda: 21} - - with handler(analysisN): - i = f1() - t = i.values(typ) - v = i.values(value) - assert t is int - assert v == 21 - - -def test_simul_analysis_apply(): - @defop - def plus1[T](x: T) -> T: - raise NotHandled - - @defop - def plus2[T](x: T) -> T: - raise NotHandled - - @defop - def times[T](x: T, y: T) -> T: - raise NotHandled - - x, y = defop(int, name="x"), defop(int, name="y") - - typ = defop(Interpretation, name="typ") - value = defop(Interpretation, name="value") - - def apply_type(op, *a, **k): - return Box(op.__type_rule__(*a, **k)) - - type_rules = {apply: apply_type} - - def plus1_value(x): - return x + 1 - - def plus2_value(x): - return plus1(plus1(x)) - - def times_value(x, y): - t = typ().value - arg = argsof(typ)[0][0].value - if t is int and arg is int: - return x * y - raise TypeError("unexpected type!") - - value_rules = { - plus1: plus1_value, - plus2: plus2_value, - times: times_value, - x: lambda: 3, - y: lambda: 4, - } - - analysisN = productN({typ: type_rules, value: value_rules}) - - def f1(): - v1 = x() # {typ: lambda: int, val: lambda: 3} - v2 = y() # {typ: lambda: int, val: lambda: 4} - v3 = plus2(v1) # {typ: lambda: int, val: lambda: 5} - v4 = times(v2, v3) # {typ: lambda: int, val: lambda: 20} - v5 = plus1(v4) # {typ: lambda: int, val: lambda: 21} - return v5 # {typ: lambda: int, val: lambda: 21} - - with handler(analysisN): - i = f1() - t = i.values(typ).value - v = i.values(value) - assert t is int - assert v == 21 - - -def test_productN_distributive(): - """Test that productN distributes over coproducts.""" - - @defop - def add[T](x: T, y: T) -> T: - raise NotHandled - - x = defop(object, name="x") - i = defop(object, name="i") - s = defop(object, name="s") - - intp1 = {add: lambda x, y: x + y} - intp2 = {x: lambda: 1} - intp3 = {x: lambda: "a"} - - term = add(x(), x()) - - prod_intp1 = productN({i: coproduct(intp2, intp1), s: coproduct(intp3, intp1)}) - prod_intp2 = coproduct( - productN({i: intp2, s: intp3}), productN({i: intp1, s: intp1}) - ) - result1 = evaluate(term, intp=prod_intp1) - result2 = evaluate(term, intp=prod_intp2) - - assert result1.values(i) == result2.values(i) == 2 - assert result1.values(s) == result2.values(s) == "aa" diff --git a/tests/test_ops_semantics.py b/tests/test_ops_semantics.py index f950949cf..17c542561 100644 --- a/tests/test_ops_semantics.py +++ b/tests/test_ops_semantics.py @@ -856,6 +856,30 @@ def Lam2[A, B]( assert actual >= {z, Lam2, add} +def test_fvsof_collection_binder(): + a, b, c, d = ( + defop(int, name="a"), + defop(int, name="b"), + defop(int, name="c"), + defop(int, name="d"), + ) + + @defop + def add(x: int, y: int) -> int: + raise NotHandled + + @defop + def let_many[A, B]( + body: Annotated[int, Scoped[A | B]], + bindings: Annotated[dict[Operation[[], int], int], Scoped[A]], + ) -> Annotated[int, Scoped[B]]: + raise NotHandled + + term = let_many(add(a(), b()), {a: c(), c: d()}) + actual = fvsof(term) + assert actual == {b, d, let_many, add} + + def test_fvsof_collection_does_not_include_apply(): x = defop(int, name="x")