Remove intp argument from apply (#333)

Author: eb8680

effectful/handlers/numpyro.py

@@ -85,7 +85,7 @@ def _to_named(a):
             return a
 
     # Convert to a term in a context that does not evaluate distribution constructors.
-    def _apply(intp, op, *args, **kwargs):
+    def _apply(op, *args, **kwargs):
         typ = op.__type_rule__(*args, **kwargs)
         if issubclass(typ, dist.Distribution):
             return defdata(op, *args, **kwargs)
@@ -133,7 +133,7 @@ def _to_positional(a, indices):
             return a
 
     # Convert to a term in a context that does not evaluate distribution constructors.
-    def _apply(intp, op, *args, **kwargs):
+    def _apply(op, *args, **kwargs):
         typ = op.__type_rule__(*args, **kwargs)
         if issubclass(typ, dist.Distribution):
             return defdata(op, *args, **kwargs)

effectful/handlers/pyro.py

@@ -369,7 +369,7 @@ def _to_named(a):
             return a
 
     # Convert to a term in a context that does not evaluate distribution constructors.
-    def _apply(intp, op, *args, **kwargs):
+    def _apply(op, *args, **kwargs):
         typ = op.__type_rule__(*args, **kwargs)
         if issubclass(
             typ, pyro.distributions.torch_distribution.TorchDistribution
@@ -417,7 +417,7 @@ def _to_positional(a, indices):
             return a
 
     # Convert to a term in a context that does not evaluate distribution constructors.
-    def _apply(intp, op, *args, **kwargs):
+    def _apply(op, *args, **kwargs):
         typ = op.__type_rule__(*args, **kwargs)
         if issubclass(
             typ, pyro.distributions.torch_distribution.TorchDistribution

effectful/handlers/torch.py

@@ -72,7 +72,7 @@ def _torch_getitem_sizeof(
 
         return defdata(torch_getitem, x, key)
 
-    def _apply(_, op, *args, **kwargs):
+    def _apply(op, *args, **kwargs):
         args, kwargs = tree.map_structure(defterm, (args, kwargs))
         return defdata(op, *args, **kwargs)
 

effectful/ops/semantics.py

@@ -11,8 +11,8 @@
 from effectful.ops.types import Expr, Interpretation, NotHandled, Operation, Term
 
 
-@defop
-def apply(intp: Interpretation, op: Operation, *args, **kwargs) -> Any:
+@defop  # type: ignore
+def apply[**P, T](op: Operation[P, T], *args: P.args, **kwargs: P.kwargs) -> T:
     """Apply ``op`` to ``args``, ``kwargs`` in interpretation ``intp``.
 
     Handling :func:`apply` changes the evaluation strategy of terms.
@@ -41,14 +41,15 @@ def apply(intp: Interpretation, op: Operation, *args, **kwargs) -> Any:
     mul(add(1, 2), 3)
 
     """
-    from effectful.internals.runtime import interpreter
+    from effectful.internals.runtime import get_interpretation
 
+    intp = get_interpretation()
     if op in intp:
-        return interpreter(intp)(intp[op])(*args, **kwargs)
+        return intp[op](*args, **kwargs)
     elif apply in intp:
-        return intp[apply](intp, op, *args, **kwargs)
+        return intp[apply](op, *args, **kwargs)
     else:
-        return op.__default_rule__(*args, **kwargs)
+        return op.__default_rule__(*args, **kwargs)  # type: ignore
 
 
 @defop  # type: ignore
@@ -208,7 +209,7 @@ def runner(intp: Interpretation):
     from effectful.internals.runtime import get_interpretation, interpreter
 
     @interpreter(get_interpretation())
-    def _reapply[**P, S](_, op: Operation[P, S], *args: P.args, **kwargs: P.kwargs):
+    def _reapply[**P, S](op: Operation[P, S], *args: P.args, **kwargs: P.kwargs):
         return op(*args, **kwargs)
 
     with interpreter({apply: _reapply, **intp}):
@@ -316,7 +317,7 @@ def typeof[T](term: Expr[T]) -> type[T]:
     """
     from effectful.internals.runtime import interpreter
 
-    with interpreter({apply: lambda _, op, *a, **k: op.__type_rule__(*a, **k)}):
+    with interpreter({apply: lambda op, *a, **k: op.__type_rule__(*a, **k)}):
         if isinstance(term, Term):
             # If term is a Term, we evaluate it to get its type
             tp = evaluate(term)
@@ -353,7 +354,7 @@ def fvsof[S](term: Expr[S]) -> set[Operation]:
 
     _fvs: set[Operation] = set()
 
-    def _update_fvs(_, op, *args, **kwargs):
+    def _update_fvs(op, *args, **kwargs):
         _fvs.add(op)
         arg_ctxs, kwarg_ctxs = op.__fvs_rule__(*args, **kwargs)
         bound_vars = set().union(

effectful/ops/syntax.py

@@ -958,10 +958,7 @@ def _(op, *args, **kwargs):
         if c:
             res = evaluate(
                 v,
-                intp={
-                    apply: lambda _, op, *a, **k: defdata(op, *a, **k),
-                    **{op: renaming[op] for op in c},
-                },
+                intp={apply: defdata, **{op: renaming[op] for op in c}},
             )
             if isinstance(i, int):
                 args_[i] = res
@@ -1067,12 +1064,7 @@ def trace[**P, T](value: Callable[P, T]) -> Callable[P, T]:
     )
     bound_sig.apply_defaults()
 
-    with interpreter(
-        {
-            apply: lambda _, op, *a, **k: defdata(op, *a, **k),
-            call: call.__default_rule__,
-        }
-    ):
+    with interpreter({apply: defdata, call: call.__default_rule__}):
         body = value(
             *[a() for a in bound_sig.args],
             **{k: v() for k, v in bound_sig.kwargs.items()},

effectful/ops/types.py

@@ -72,10 +72,9 @@ def __fvs_rule__(
 
     @typing.final
     def __call__(self, *args: Q.args, **kwargs: Q.kwargs) -> V:
-        from effectful.internals.runtime import get_interpretation
         from effectful.ops.semantics import apply
 
-        return apply.__default_rule__(get_interpretation(), self, *args, **kwargs)  # type: ignore
+        return apply.__default_rule__(self, *args, **kwargs)  # type: ignore
 
     def __repr__(self):
         return f"{self.__class__.__name__}({self.__name__}, {self.__signature__})"
@@ -146,7 +145,7 @@ def term_str(term):
                 )
             return str(term)
 
-        def _apply(_, op, *args, **kwargs) -> str:
+        def _apply(op, *args, **kwargs) -> str:
             args_str = ", ".join(map(term_str, args)) if args else ""
             kwargs_str = (
                 ", ".join(f"{k}={term_str(v)}" for k, v in kwargs.items())

tests/test_ops_semantics.py

@@ -8,7 +8,6 @@
 
 import effectful.handlers.numbers  # noqa: F401
 from effectful.ops.semantics import (
-    apply,
     coproduct,
     evaluate,
     fvsof,
@@ -578,7 +577,6 @@ def g(x: str, y: bool) -> str:
     runner(i)
     product(i, i)
     coproduct(i, i)
-    apply(i, f, 1)
     evaluate(0, intp=i)
 
     # include tests with inlined interpretation, because mypy might do inference
@@ -593,7 +591,6 @@ def g(x: str, y: bool) -> str:
         {f: lambda x: x + 1, g: lambda x, y: x + str(y)},
         {f: lambda x: x + 1, g: lambda x, y: x + str(y)},
     )
-    apply({f: lambda x: x + 1, g: lambda x, y: x + str(y)}, f, 1)
     evaluate(0, intp={f: lambda x: x + 1, g: lambda x, y: x + str(y)})