effectful 0.0.1__py3-none-any.whl → 0.1.0__py3-none-any.whl

effectful/handlers/indexed.py

@@ -6,7 +6,7 @@ import torch
 
 from effectful.handlers.torch import Indexable, sizesof
 from effectful.ops.syntax import deffn, defop
-from effectful.ops.types import Operation, Term
+from effectful.ops.types import Operation
 
 K = TypeVar("K")
 T = TypeVar("T")
@@ -61,16 +61,6 @@ class IndexSet(Dict[str, Set[int]]):
     def __hash__(self):
         return hash(frozenset((k, frozenset(vs)) for k, vs in self.items()))
 
-    def _to_handler(self):
-        """Return an effectful handler that binds each index variable to a
-        tensor of its possible index values.
-
-        """
-        return {
-            name_to_sym(k): functools.partial(lambda v: v, torch.tensor(list(v)))
-            for k, v in self.items()
-        }
-
 
 def union(*indexsets: IndexSet) -> IndexSet:
     """
@@ -166,17 +156,9 @@ def indices_of(value: Any) -> IndexSet:
     :param kwargs: Additional keyword arguments used by specific implementations.
     :return: A :class:`IndexSet` containing the indices on which the value is supported.
     """
-    if isinstance(value, Term):
-        return IndexSet(
-            **{
-                k.__name__: set(range(v))  # type:ignore
-                for (k, v) in sizesof(value).items()
-            }
-        )
-    elif isinstance(value, torch.distributions.Distribution):
-        return indices_of(value.sample())
-
-    return IndexSet()
+    return IndexSet(
+        **{getattr(k, "__name__"): set(range(v)) for (k, v) in sizesof(value).items()}
+    )
 
 
 @functools.lru_cache(maxsize=None)

effectful/handlers/numbers.py

@@ -1,11 +1,15 @@
+"""
+This module provides a term representation for numbers and operations on them.
+"""
+
 import numbers
 import operator
 from typing import Any, TypeVar
 
 from typing_extensions import ParamSpec
 
-from effectful.ops.syntax import NoDefaultRule, defdata, defop, syntactic_eq
-from effectful.ops.types import Operation, Term
+from effectful.ops.syntax import defdata, defop, syntactic_eq
+from effectful.ops.types import Expr, Operation, Term
 
 P = ParamSpec("P")
 Q = ParamSpec("Q")
@@ -20,7 +24,7 @@ T_Number = TypeVar("T_Number", bound=numbers.Number)
 @numbers.Number.register
 class _NumberTerm(Term[numbers.Number]):
     def __init__(
-        self, op: Operation[..., numbers.Number], args: tuple, kwargs: dict
+        self, op: Operation[..., numbers.Number], *args: Expr, **kwargs: Expr
     ) -> None:
         self._op = op
         self._args = args
@@ -56,7 +60,7 @@ def _wrap_cmp(op):
         if not any(isinstance(a, Term) for a in (x, y)):
             return op(x, y)
         else:
-            raise NoDefaultRule
+            raise NotImplementedError
 
     _wrapped_op.__name__ = op.__name__
     return _wrapped_op
@@ -67,7 +71,7 @@ def _wrap_binop(op):
         if not any(isinstance(a, Term) for a in (x, y)):
             return op(x, y)
         else:
-            raise NoDefaultRule
+            raise NotImplementedError
 
     _wrapped_op.__name__ = op.__name__
     return _wrapped_op
@@ -78,7 +82,7 @@ def _wrap_unop(op):
         if not isinstance(x, Term):
             return op(x)
         else:
-            raise NoDefaultRule
+            raise NotImplementedError
 
     _wrapped_op.__name__ = op.__name__
     return _wrapped_op

effectful/handlers/pyro.py

@@ -264,7 +264,7 @@ class PositionalDistribution(pyro.distributions.torch_distribution.TorchDistribu
         self, base_dist: pyro.distributions.torch_distribution.TorchDistribution
     ):
         self.base_dist = base_dist
-        self.indices = sizesof(base_dist.sample())
+        self.indices = sizesof(base_dist)
 
         n_base = len(base_dist.batch_shape) + len(base_dist.event_shape)
         self.naming = Naming.from_shape(self.indices.keys(), n_base)
@@ -361,7 +361,7 @@ class NamedDistribution(pyro.distributions.torch_distribution.TorchDistribution)
         self.names = names
 
         assert 1 <= len(names) <= len(base_dist.batch_shape)
-        base_indices = sizesof(base_dist.sample())
+        base_indices = sizesof(base_dist)
         assert not any(n in base_indices for n in names)
 
         n_base = len(base_dist.batch_shape) + len(base_dist.event_shape)

effectful/handlers/torch.py

@@ -12,10 +12,9 @@ import tree
 from typing_extensions import ParamSpec
 
 import effectful.handlers.numbers  # noqa: F401
-from effectful.internals.base_impl import _BaseTerm
 from effectful.internals.runtime import interpreter
 from effectful.ops.semantics import apply, evaluate, fvsof, typeof
-from effectful.ops.syntax import NoDefaultRule, defdata, defop
+from effectful.ops.syntax import defdata, defop
 from effectful.ops.types import Expr, Operation, Term
 
 P = ParamSpec("P")
@@ -90,6 +89,11 @@ def sizesof(value: Expr) -> Mapping[Operation[[], int], int]:
     >>> sizesof(Indexable(torch.ones(2, 3))[a(), b()])
     {a: 2, b: 3}
     """
+    if isinstance(value, torch.distributions.Distribution) and not isinstance(
+        value, Term
+    ):
+        return {v: s for a in value.__dict__.values() for v, s in sizesof(a).items()}
+
     sizes: dict[Operation[[], int], int] = {}
 
     def _torch_getitem_sizeof(
@@ -111,12 +115,12 @@ def sizesof(value: Expr) -> Mapping[Operation[[], int], int]:
                         )
                     sizes[k.op] = shape[i]
 
-        return torch_getitem.__free_rule__(x, key)
+        return defdata(torch_getitem, x, key)
 
     with interpreter(
         {
             torch_getitem: _torch_getitem_sizeof,
-            apply: lambda _, op, *a, **k: op.__free_rule__(*a, **k),
+            apply: lambda _, op, *a, **k: defdata(op, *a, **k),
         }
     ):
         evaluate(value)
@@ -204,7 +208,7 @@ def _register_torch_op(torch_fn: Callable[P, T]):
     @defop
     def _torch_op(*args, **kwargs) -> torch.Tensor:
 
-        tm = _torch_op.__free_rule__(*args, **kwargs)
+        tm = defdata(_torch_op, *args, **kwargs)
         sized_fvs = sizesof(tm)
 
         if (
@@ -214,7 +218,7 @@ def _register_torch_op(torch_fn: Callable[P, T]):
             and args[1]
             and all(isinstance(k, Term) and k.op in sized_fvs for k in args[1])
         ):
-            raise NoDefaultRule
+            raise NotImplementedError
         elif sized_fvs and set(sized_fvs.keys()) == fvsof(tm) - {
             torch_getitem,
             _torch_op,
@@ -230,7 +234,7 @@ def _register_torch_op(torch_fn: Callable[P, T]):
         ):
             return typing.cast(torch.Tensor, torch_fn(*args, **kwargs))
         else:
-            raise NoDefaultRule
+            raise NotImplementedError
 
     functools.update_wrapper(_torch_op, torch_fn)
     return _torch_op
@@ -315,7 +319,7 @@ class Indexable:
 
 
 @defdata.register(torch.Tensor)
-def _embed_tensor(op, args, kwargs):
+def _embed_tensor(op, *args, **kwargs):
     if (
         op is torch_getitem
         and not isinstance(args[0], Term)
@@ -328,10 +332,29 @@ def _embed_tensor(op, args, kwargs):
     ):
         return _EagerTensorTerm(args[0], args[1])
     else:
-        return _TensorTerm(op, args, kwargs)
+        return _TensorTerm(op, *args, **kwargs)
+
 
+class _TensorTerm(Term[torch.Tensor]):
+    def __init__(
+        self, op: Operation[..., torch.Tensor], *args: Expr, **kwargs: Expr
+    ) -> None:
+        self._op = op
+        self._args = args
+        self._kwargs = kwargs
+
+    @property
+    def op(self) -> Operation[..., torch.Tensor]:
+        return self._op
+
+    @property
+    def args(self) -> tuple:
+        return self._args
+
+    @property
+    def kwargs(self) -> dict:
+        return self._kwargs
 
-class _TensorTerm(_BaseTerm[torch.Tensor]):
     def __getitem__(
         self, key: Union[Expr[IndexElement], Tuple[Expr[IndexElement], ...]]
     ) -> Expr[torch.Tensor]:

effectful/ops/semantics.py

@@ -5,7 +5,7 @@ from typing import Any, Callable, Optional, Set, Type, TypeVar
 import tree
 from typing_extensions import ParamSpec
 
-from effectful.ops.syntax import NoDefaultRule, deffn, defop, defterm
+from effectful.ops.syntax import deffn, defop
 from effectful.ops.types import Expr, Interpretation, Operation, Term
 
 P = ParamSpec("P")
@@ -15,10 +15,8 @@ T = TypeVar("T")
 V = TypeVar("V")
 
 
-@defop  # type: ignore
-def apply(
-    intp: Interpretation[S, T], op: Operation[P, S], *args: P.args, **kwargs: P.kwargs
-) -> T:
+@defop
+def apply(intp: Interpretation, op: Operation, *args, **kwargs) -> Any:
     """Apply ``op`` to ``args``, ``kwargs`` in interpretation ``intp``.
 
     Handling :func:`apply` changes the evaluation strategy of terms.
@@ -50,7 +48,7 @@ def apply(
     elif apply in intp:
         return intp[apply](intp, op, *args, **kwargs)
     else:
-        return op.__default_rule__(*args, **kwargs)  # type: ignore
+        return op.__default_rule__(*args, **kwargs)
 
 
 @defop  # type: ignore
@@ -60,9 +58,6 @@ def call(fn: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> T:
     This operation is invoked by the ``__call__`` method of a callable term.
 
     """
-    if not isinstance(fn, Term):
-        fn = defterm(fn)
-
     if isinstance(fn, Term) and fn.op is deffn:
         body: Expr[Callable[P, T]] = fn.args[0]
         argvars: tuple[Operation, ...] = fn.args[1:]
@@ -73,8 +68,10 @@ def call(fn: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> T:
         }
         with handler(subs):
             return evaluate(body)
+    elif not any(isinstance(a, Term) for a in tree.flatten((fn, args, kwargs))):
+        return fn(*args, **kwargs)
     else:
-        raise NoDefaultRule
+        raise NotImplementedError
 
 
 @defop
@@ -93,9 +90,7 @@ def fwd(*args, **kwargs) -> Any:
     raise RuntimeError("fwd should only be called in the context of a handler")
 
 
-def coproduct(
-    intp: Interpretation[S, T], intp2: Interpretation[S, T]
-) -> Interpretation[S, T]:
+def coproduct(intp: Interpretation, intp2: Interpretation) -> Interpretation:
     """The coproduct of two interpretations handles any effect that is handled
     by either. If both interpretations handle an effect, ``intp2`` takes
     precedence.
@@ -151,7 +146,7 @@ def coproduct(
         if op is fwd or op is _get_args:
             res[op] = i2  # fast path for special cases, should be equivalent if removed
         else:
-            i1 = intp.get(op, op.__default_rule__)  # type: ignore
+            i1 = intp.get(op, op.__default_rule__)
 
             # calling fwd in the right handler should dispatch to the left handler
             res[op] = _set_prompt(fwd, _restore_args(_save_args(i1)), _save_args(i2))
@@ -159,9 +154,7 @@ def coproduct(
     return res
 
 
-def product(
-    intp: Interpretation[S, T], intp2: Interpretation[S, T]
-) -> Interpretation[S, T]:
+def product(intp: Interpretation, intp2: Interpretation) -> Interpretation:
     """The product of two interpretations handles any effect that is handled by
     ``intp2``. Handlers in ``intp2`` may override handlers in ``intp``, but
     those changes are not visible to the handlers in ``intp``. In this way,
@@ -207,7 +200,7 @@ def product(
 
 
 @contextlib.contextmanager
-def runner(intp: Interpretation[S, T]):
+def runner(intp: Interpretation):
     """Install an interpretation by taking a product with the current
     interpretation.
 
@@ -223,7 +216,7 @@ def runner(intp: Interpretation[S, T]):
 
 
 @contextlib.contextmanager
-def handler(intp: Interpretation[S, T]):
+def handler(intp: Interpretation):
     """Install an interpretation by taking a coproduct with the current
     interpretation.
 
@@ -234,7 +227,7 @@ def handler(intp: Interpretation[S, T]):
         yield intp
 
 
-def evaluate(expr: Expr[T], *, intp: Optional[Interpretation[S, T]] = None) -> Expr[T]:
+def evaluate(expr: Expr[T], *, intp: Optional[Interpretation] = None) -> Expr[T]:
     """Evaluate expression ``expr`` using interpretation ``intp``. If no
     interpretation is provided, uses the current interpretation.
 
@@ -245,7 +238,7 @@ def evaluate(expr: Expr[T], *, intp: Optional[Interpretation[S, T]] = None) -> E
 
     >>> @defop
     ... def add(x: int, y: int) -> int:
-    ...     raise NoDefaultRule
+    ...     raise NotImplementedError
     >>> expr = add(1, add(2, 3))
     >>> expr
     add(1, add(2, 3))
@@ -258,13 +251,11 @@ def evaluate(expr: Expr[T], *, intp: Optional[Interpretation[S, T]] = None) -> E
 
         intp = get_interpretation()
 
-    expr = defterm(expr) if not isinstance(expr, Term) else expr
-
     if isinstance(expr, Term):
         (args, kwargs) = tree.map_structure(
             functools.partial(evaluate, intp=intp), (expr.args, expr.kwargs)
         )
-        return apply.__default_rule__(intp, expr.op, *args, **kwargs)  # type: ignore
+        return apply.__default_rule__(intp, expr.op, *args, **kwargs)
     elif tree.is_nested(expr):
         return tree.map_structure(functools.partial(evaluate, intp=intp), expr)
     else:
@@ -280,7 +271,7 @@ def typeof(term: Expr[T]) -> Type[T]:
 
     >>> @defop
     ... def cmp(x: int, y: int) -> bool:
-    ...     raise NoDefaultRule
+    ...     raise NotImplementedError
     >>> typeof(cmp(1, 2))
     <class 'bool'>
 
@@ -290,7 +281,7 @@ def typeof(term: Expr[T]) -> Type[T]:
     >>> T = TypeVar('T')
     >>> @defop
     ... def if_then_else(x: bool, a: T, b: T) -> T:
-    ...     raise NoDefaultRule
+    ...     raise NotImplementedError
     >>> typeof(if_then_else(True, 0, 1))
     <class 'int'>
 
@@ -298,7 +289,7 @@ def typeof(term: Expr[T]) -> Type[T]:
     from effectful.internals.runtime import interpreter
 
     with interpreter({apply: lambda _, op, *a, **k: op.__type_rule__(*a, **k)}):
-        return evaluate(term)  # type: ignore
+        return evaluate(term) if isinstance(term, Term) else type(term)  # type: ignore
 
 
 def fvsof(term: Expr[S]) -> Set[Operation]:
@@ -308,7 +299,7 @@ def fvsof(term: Expr[S]) -> Set[Operation]:
 
     >>> @defop
     ... def f(x: int, y: int) -> int:
-    ...     raise NoDefaultRule
+    ...     raise NotImplementedError
     >>> fvsof(f(1, 2))
     {f}
 
@@ -319,7 +310,12 @@ def fvsof(term: Expr[S]) -> Set[Operation]:
 
     def _update_fvs(_, op, *args, **kwargs):
         _fvs.add(op)
-        for bound_var in op.__fvs_rule__(*args, **kwargs):
+        arg_ctxs, kwarg_ctxs = op.__fvs_rule__(*args, **kwargs)
+        bound_vars = set().union(
+            *(a for a in arg_ctxs),
+            *(k for k in kwarg_ctxs.values()),
+        )
+        for bound_var in bound_vars:
             if bound_var in _fvs:
                 _fvs.remove(bound_var)