Allow _BaseOperation subclasses to have an overrideable apply method (#414)

* stash

* fixes

* initial

* wip

* lint

* ensure each subclass has a fresh operation

* wip

* wip

* lint

* wip

* wip

* lint

* refactor class method support

* move defops

* fix test

* remove singledispatch case and add test

* move definition

* cleanup

* simplify

* cleanup

* lint

* fix failing test

* fix classmethod

* __isabstractmethod__

* revert

---------

Co-authored-by: Eli <eli@basis.ai>

Author: jfeser

docs/source/lambda_.py

@@ -102,7 +102,7 @@ def sort_add(x: Expr[int], y: Expr[int]) -> Expr[int]:
         case Term(add_, (a, Term(vx, ()))), Term(vy, ()) if add_ == add and id(vx) > id(
             vy
         ):
-            return (a + vy()) + vx()  # type: ignore
+            return (a + vy()) + vx()
         case _:
             return fwd()
 

docs/source/semi_ring.py

@@ -174,11 +174,11 @@ def vertical_fusion[S, T](e1: T, x: Operation[[], T], e2: S) -> S:
             return evaluate(
                 Sum(
                     e_sum,  # type: ignore
-                    k1,  # type: ignore
-                    v1,  # type: ignore
+                    k1,
+                    v1,
                     Let(
                         e_lhs,
-                        v2,  # type: ignore
+                        v2,
                         Let(k1(), k2, Dict(k2(), Let(e_lhs, k2, e_rhs))),  # type: ignore
                     ),
                 )

effectful/handlers/numpyro.py

@@ -242,22 +242,22 @@ def args(self):
     def kwargs(self):
         return self._kwargs
 
-    @defop  # type: ignore
     @property
+    @defop
     def batch_shape(self) -> tuple[int, ...]:
         if not (self._is_eager):
             raise NotHandled
         return self._pos_base_dist.batch_shape[len(self._indices) :]
 
-    @defop  # type: ignore
     @property
+    @defop
     def has_rsample(self) -> bool:
         if not (self._is_eager):
             raise NotHandled
         return self._pos_base_dist.has_rsample
 
-    @defop  # type: ignore
     @property
+    @defop
     def event_shape(self) -> tuple[int, ...]:
         if not (self._is_eager):
             raise NotHandled
@@ -312,8 +312,8 @@ def log_prob(self, value) -> jax.Array:
         ind_log_prob = self._reindex_sample(pos_log_prob, sample_shape)
         return ind_log_prob
 
-    @defop  # type: ignore
     @property
+    @defop
     def mean(self) -> jax.Array:
         if not self._is_eager:
             raise NotHandled
@@ -322,8 +322,8 @@ def mean(self) -> jax.Array:
         except NotImplementedError:
             raise RuntimeError(f"mean is not implemented for {type(self).__name__}")
 
-    @defop  # type: ignore
     @property
+    @defop
     def variance(self) -> jax.Array:
         if not self._is_eager:
             raise NotHandled

effectful/internals/runtime.py

@@ -40,7 +40,7 @@ def _get_args() -> tuple[tuple, Mapping]:
 def _restore_args[**P, T](fn: Callable[P, T]) -> Callable[P, T]:
     @functools.wraps(fn)
     def _cont_wrapper(*a: P.args, **k: P.kwargs) -> T:
-        a, k = (a, k) if a or k else _get_args()  # type: ignore
+        a, k = (a, k) if a or k else _get_args()
         return fn(*a, **k)
 
     return _cont_wrapper

effectful/ops/semantics.py

@@ -16,7 +16,7 @@
 )
 
 
-@defop  # type: ignore
+@defop
 def apply[**P, T](op: Operation[P, T], *args: P.args, **kwargs: P.kwargs) -> T:
     """Apply ``op`` to ``args``, ``kwargs`` in interpretation ``intp``.
 
@@ -38,23 +38,14 @@ def apply[**P, T](op: Operation[P, T], *args: P.args, **kwargs: P.kwargs) -> T:
 
     By installing an :func:`apply` handler, we capture the term instead:
 
-    >>> def default(*args, **kwargs):
-    ...     raise NotHandled
-    >>> with handler({apply: default }):
+    >>> from effectful.ops.syntax import defdata
+    >>> with handler({apply: defdata}):
     ...     term = mul(add(1, 2), 3)
     >>> print(str(term))
     mul(add(1, 2), 3)
 
     """
-    from effectful.internals.runtime import get_interpretation
-
-    intp = get_interpretation()
-    if op in intp:
-        return intp[op](*args, **kwargs)
-    elif apply in intp:
-        return intp[apply](op, *args, **kwargs)
-    else:
-        return op.__default_rule__(*args, **kwargs)  # type: ignore
+    return op.__default_rule__(*args, **kwargs)  # type: ignore
 
 
 @defop