Implement MeasurableElemwise

Subclasses are measurable by default, and they must specify what scalar op types are compatible.

Author: Ricardo Vieira

aeppl/abstract.py

@@ -1,11 +1,12 @@
 import abc
 from copy import copy
 from functools import singledispatch
-from typing import Callable, List
+from typing import Callable, List, Tuple
 
 from aesara.graph.basic import Apply, Variable
 from aesara.graph.op import Op
 from aesara.graph.utils import MetaType
+from aesara.tensor.elemwise import Elemwise
 from aesara.tensor.random.op import RandomVariable
 
 
@@ -116,3 +117,20 @@ def assign_custom_measurable_outputs(
     _get_measurable_outputs.register(new_op_type)(measurable_outputs_fn)
 
     return new_node
+
+
+class MeasurableElemwise(Elemwise):
+    """Base class for Measurable Elemwise variables"""
+
+    valid_scalar_types: Tuple[MetaType, ...] = ()
+
+    def __init__(self, scalar_op, *args, **kwargs):
+        if not isinstance(scalar_op, self.valid_scalar_types):
+            raise TypeError(
+                f"scalar_op {scalar_op} is not valid for class {self.__class__}. "
+                f"Acceptable types are {self.valid_scalar_types}"
+            )
+        super().__init__(scalar_op, *args, **kwargs)
+
+
+MeasurableVariable.register(MeasurableElemwise)

aeppl/censoring.py

@@ -10,16 +10,22 @@
 from aesara.tensor.elemwise import Elemwise
 from aesara.tensor.var import TensorConstant
 
-from aeppl.abstract import MeasurableVariable, assign_custom_measurable_outputs
+from aeppl.abstract import (
+    MeasurableElemwise,
+    MeasurableVariable,
+    assign_custom_measurable_outputs,
+)
 from aeppl.logprob import CheckParameterValue, _logcdf, _logprob
 from aeppl.rewriting import measurable_ir_rewrites_db
 
 
-class MeasurableClip(Elemwise):
+class MeasurableClip(MeasurableElemwise):
     """A placeholder used to specify a log-likelihood for a clipped RV sub-graph."""
 
+    valid_scalar_types = (Clip,)
+
 
-MeasurableVariable.register(MeasurableClip)
+measurable_clip = MeasurableClip(scalar_clip)
 
 
 @node_rewriter(tracks=[Elemwise])
@@ -54,10 +60,9 @@ def find_measurable_clips(
     lower_bound = lower_bound if (lower_bound is not base_var) else at.constant(-np.inf)
     upper_bound = upper_bound if (upper_bound is not base_var) else at.constant(np.inf)
 
-    clipped_op = MeasurableClip(scalar_clip)
     # Make base_var unmeasurable
     unmeasurable_base_var = assign_custom_measurable_outputs(base_var.owner)
-    clipped_rv_node = clipped_op.make_node(
+    clipped_rv_node = measurable_clip.make_node(
         unmeasurable_base_var, lower_bound, upper_bound
     )
     clipped_rv = clipped_rv_node.outputs[0]

aeppl/transforms.py

@@ -20,6 +20,7 @@
 from aesara.tensor.var import TensorVariable
 
 from aeppl.abstract import (
+    MeasurableElemwise,
     MeasurableVariable,
     _get_measurable_outputs,
     assign_custom_measurable_outputs,
@@ -213,9 +214,11 @@ def apply(self, fgraph: FunctionGraph):
         return self.default_transform_rewrite.rewrite(fgraph)
 
 
-class MeasurableTransform(Elemwise):
+class MeasurableTransform(MeasurableElemwise):
     """A placeholder used to specify a log-likelihood for a transformed measurable variable"""
 
+    valid_scalar_types = (Exp, Log, Add, Mul)
+
     # Cannot use `transform` as name because it would clash with the property added by
     # the `TransformValuesRewrite`
     transform_elemwise: RVTransform
@@ -229,9 +232,6 @@ def __init__(
         super().__init__(*args, **kwargs)
 
 
-MeasurableVariable.register(MeasurableTransform)
-
-
 @_get_measurable_outputs.register(MeasurableTransform)
 def _get_measurable_outputs_Transform(op, node):
     return [node.default_output()]
@@ -261,7 +261,7 @@ def find_measurable_transforms(
 ) -> Optional[List[Node]]:
     """Find measurable transformations from Elemwise operators."""
     scalar_op = node.op.scalar_op
-    if not isinstance(scalar_op, (Exp, Log, Add, Mul)):
+    if not isinstance(scalar_op, MeasurableTransform.valid_scalar_types):
         return None
 
     # Node was already converted

tests/test_abstract.py

@@ -1,8 +1,13 @@
+import re
+
 import aesara.tensor as at
 import pytest
+from aesara.scalar import Exp, exp
 from aesara.tensor.random.basic import NormalRV
 
 from aeppl.abstract import (
+    MeasurableElemwise,
+    MeasurableVariable,
     UnmeasurableVariable,
     _get_measurable_outputs,
     assign_custom_measurable_outputs,
@@ -92,3 +97,16 @@ def test_assign_custom_measurable_outputs():
 
     with pytest.raises(ValueError):
         assign_custom_measurable_outputs(unmeas_X_rv.owner, lambda x: x)
+
+
+def test_measurable_elemwise():
+    # Default does not accept any scalar_op
+    with pytest.raises(TypeError, match=re.escape("scalar_op exp is not valid")):
+        MeasurableElemwise(exp)
+
+    class TestMeasurableElemwise(MeasurableElemwise):
+        valid_scalar_types = (Exp,)
+
+    measurable_exp_op = TestMeasurableElemwise(scalar_op=exp)
+    measurable_exp = measurable_exp_op(0.0)
+    assert isinstance(measurable_exp.owner.op, MeasurableVariable)