effectful 0.2.1__py3-none-any.whl → 0.2.2__py3-none-any.whl

effectful/handlers/indexed.py

@@ -233,7 +233,8 @@ def gather(value: torch.Tensor, indexset: IndexSet) -> torch.Tensor:
         if k in indexset_vars
     }
 
-    return deffn(value, *binding.keys())(*[v() for v in binding.values()])
+    args = [v() for v in binding.values()]
+    return deffn(value, *binding.keys())(*args)
 
 
 def stack(

effectful/internals/product_n.py

@@ -0,0 +1,163 @@
+import dataclasses
+import functools
+from collections.abc import Callable, Mapping
+from typing import Any
+
+import tree
+
+from effectful.ops.semantics import apply, coproduct, handler
+from effectful.ops.syntax import defop
+from effectful.ops.types import Interpretation, 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.")
+
+
+@dataclasses.dataclass
+class Product[S, T]:
+    values: object
+
+
+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 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 = tree.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

effectful/ops/semantics.py

@@ -282,6 +282,21 @@ def evaluate[T](expr: Expr[T], *, intp: Interpretation | None = None) -> Expr[T]
         return typing.cast(T, expr)
 
 
+def _simple_type(tp: type) -> type:
+    """Convert a type object into a type that can be dispatched on."""
+    if isinstance(tp, typing.TypeVar):
+        tp = (
+            tp.__bound__
+            if tp.__bound__
+            else tp.__constraints__[0]
+            if tp.__constraints__
+            else object
+        )
+    if isinstance(tp, types.UnionType):
+        raise TypeError(f"Union types are not supported: {tp}")
+    return typing.get_origin(tp) or tp
+
+
 def typeof[T](term: Expr[T]) -> type[T]:
     """Return the type of an expression.
 
@@ -310,19 +325,7 @@ def typeof[T](term: Expr[T]) -> type[T]:
         if isinstance(term, Term):
             # If term is a Term, we evaluate it to get its type
             tp = evaluate(term)
-            if isinstance(tp, typing.TypeVar):
-                tp = (
-                    tp.__bound__
-                    if tp.__bound__
-                    else tp.__constraints__[0]
-                    if tp.__constraints__
-                    else object
-                )
-            if isinstance(tp, types.UnionType):
-                raise TypeError(
-                    f"Cannot determine type of {term} because it is a union type: {tp}"
-                )
-            return typing.get_origin(tp) or tp  # type: ignore
+            return _simple_type(typing.cast(type, tp))
         else:
             return type(term)
 

effectful/ops/syntax.py

@@ -924,8 +924,12 @@ def defdata[T](
     When an Operation whose return type is `Callable` is passed to :func:`defdata`,
     it is reconstructed as a :class:`_CallableTerm`, which implements the :func:`__call__` method.
     """
-    from effectful.ops.semantics import apply, evaluate, typeof
+    from effectful.internals.product_n import productN
+    from effectful.internals.runtime import interpreter
+    from effectful.ops.semantics import _simple_type, apply, evaluate
 
+    # If this operation binds variables, we need to rename them in the
+    # appropriate parts of the child term.
     bindings: inspect.BoundArguments = op.__fvs_rule__(*args, **kwargs)
     renaming = {
         var: defop(var)
@@ -933,24 +937,54 @@ def defdata[T](
         for var in bound_vars
     }
 
-    renamed_args: inspect.BoundArguments = op.__signature__.bind(*args, **kwargs)
+    # Analysis for type computation and term reconstruction
+    typ = defop(object, name="typ")
+    cast = defop(object, name="cast")
+
+    def apply_type(op, *args, **kwargs):
+        assert isinstance(op, Operation)
+        tp = op.__type_rule__(*args, **kwargs)
+        return tp
+
+    def apply_cast(op, *args, **kwargs):
+        assert isinstance(op, Operation)
+        full_type = typ()
+        dispatch_type = _simple_type(full_type)
+        return __dispatch(dispatch_type)(op, *args, **kwargs)
+
+    analysis = productN({typ: {apply: apply_type}, cast: {apply: apply_cast}})
+
+    def evaluate_with_renaming(expr, ctx):
+        """Evaluate an expression with renaming applied."""
+        renaming_ctx = {
+            old_var: new_var for old_var, new_var in renaming.items() if old_var in ctx
+        }
+
+        # Note: coproduct cannot be used to compose these interpretations
+        # because evaluate will only do operation replacement when the handler
+        # is operation typed, which coproduct does not satisfy.
+        with interpreter(analysis | renaming_ctx):
+            result = evaluate(expr)
+
+        return result
+
+    renamed_args = op.__signature__.bind(*args, **kwargs)
     renamed_args.apply_defaults()
 
     args_ = [
-        evaluate(
-            arg, intp={apply: defdata, **{v: renaming[v] for v in bindings.args[i]}}
-        )
-        for i, arg in enumerate(renamed_args.args)
+        evaluate_with_renaming(arg, bindings.args[i])
+        for (i, arg) in enumerate(renamed_args.args)
     ]
     kwargs_ = {
-        k: evaluate(
-            arg, intp={apply: defdata, **{v: renaming[v] for v in bindings.kwargs[k]}}
-        )
-        for k, arg in renamed_args.kwargs.items()
+        k: evaluate_with_renaming(v, bindings.kwargs[k])
+        for (k, v) in renamed_args.kwargs.items()
     }
 
-    base_term = __dispatch(typing.cast(type[T], object))(op, *args_, **kwargs_)
-    return __dispatch(typeof(base_term))(op, *args_, **kwargs_)
+    # Build the final term with type analysis
+    with interpreter(analysis):
+        result = op(*args_, **kwargs_)
+
+    return result.values(cast)  # type: ignore
 
 
 @defterm.register(object)

{effectful-0.2.1.dist-info → effectful-0.2.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: effectful
-Version: 0.2.1
+Version: 0.2.2
 Summary: Metaprogramming infrastructure
 Author: Basis
 License-Expression: Apache-2.0
@@ -15,7 +15,7 @@ Classifier: Operating System :: POSIX :: Linux
 Classifier: Operating System :: MacOS :: MacOS X
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
-Requires-Python: >=3.12
+Requires-Python: <3.14,>=3.12
 Description-Content-Type: text/x-rst
 License-File: LICENSE.md
 Provides-Extra: torch

{effectful-0.2.1.dist-info → effectful-0.2.2.dist-info}/RECORD

@@ -1,7 +1,7 @@
 effectful/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 effectful/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 effectful/handlers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-effectful/handlers/indexed.py,sha256=ZY-8w32a1PKGVScwXjbgByI3wRHvfxuuXJVwLlp0rgw,12622
+effectful/handlers/indexed.py,sha256=TL3Y9EZ9Q4b12qtIEg1iWZ-mk97rZHw3zY9XAXKXCCM,12638
 effectful/handlers/numpyro.py,sha256=RWoBNpHLr5KdotN5Vu118jY0kn_p6NaejcnJgZJLehw,19457
 effectful/handlers/pyro.py,sha256=qVl1wson02pyV8YHGf93KDnYEp5pGmhKEwji95OYBl8,26486
 effectful/handlers/torch.py,sha256=NNM7mxqZskEBCjsl25kHI95WlXG9aeD7FaSkXkoLZ_I,24330
@@ -12,15 +12,16 @@ effectful/handlers/jax/numpy/__init__.py,sha256=Kmvya0QI-GA56pPf1as-wYOuZFngOBLt
 effectful/handlers/jax/numpy/linalg.py,sha256=9DiaYYG4SztmO-VkmMH3dVvULtMK-zEgbV9oNQFkFo8,350
 effectful/handlers/jax/scipy/special.py,sha256=yTIECFtQVPgraonrPlyenjvcnEYchZwIZC-5CSkF-lA,299
 effectful/internals/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+effectful/internals/product_n.py,sha256=Ii06tiGvuvMAzv_wUDBnjUGPB1ZdVxhJPj2uqBeAJyQ,5842
 effectful/internals/runtime.py,sha256=aLWol7sR1yHekn7zNz1evHKHARjiT1tnkmByLHPHBGc,1811
 effectful/internals/tensor_utils.py,sha256=3QCSUqdxCXod3dsY3oRMcg36Rqr8pVX-ktEyCEkeODo,1173
 effectful/internals/unification.py,sha256=MD0beZ29by9hbJaowafBLIZuKPF5MxHXezuGpSG_wCA,30254
 effectful/ops/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-effectful/ops/semantics.py,sha256=u6i7RTpiv5fCnOdh_pOL-5CSsO4QtW-SjfGEQiizeWk,11922
-effectful/ops/syntax.py,sha256=PE9iuLIzwqFkkk99bBZWk93X4sqSUhtRPePiTQRnAdk,55825
+effectful/ops/semantics.py,sha256=9fvBwER-ZF940b28ZS1fKvxl7ra6GafBg9FTIw2IR4c,11915
+effectful/ops/syntax.py,sha256=lP0_vX8Q_B6XD2OuDoMxBMV_fQIAOLy0QtFYtKV74Mc,57072
 effectful/ops/types.py,sha256=W1gZJaBnX7_nFpWrG3vfCBQPSun3Gc9PqT61ls8B3EA,6599
-effectful-0.2.1.dist-info/licenses/LICENSE.md,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-effectful-0.2.1.dist-info/METADATA,sha256=AgFY63eWiiFpOlvxgWa-3GOV9eSy__N0gF6ZIs5ljdc,5300
-effectful-0.2.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-effectful-0.2.1.dist-info/top_level.txt,sha256=gtuJfrE2nXil_lZLCnqWF2KAbOnJs9ILNvK8WnkRzbs,10
-effectful-0.2.1.dist-info/RECORD,,
+effectful-0.2.2.dist-info/licenses/LICENSE.md,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+effectful-0.2.2.dist-info/METADATA,sha256=9jZF6agI9xsyZ0ZUXYfSMH6ku4P-IamE_A-I9oi-T4U,5306
+effectful-0.2.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+effectful-0.2.2.dist-info/top_level.txt,sha256=gtuJfrE2nXil_lZLCnqWF2KAbOnJs9ILNvK8WnkRzbs,10
+effectful-0.2.2.dist-info/RECORD,,