Rename MeasurableVariable to MeasurableOp

Also:
* Introduce MeasurableOpMixin for string representation
* Subclass directly instead of registering manually

Author: ricardoV94

pymc/distributions/distribution.py

@@ -50,7 +50,7 @@
     rv_size_is_none,
     shape_from_dims,
 )
-from pymc.logprob.abstract import MeasurableVariable, _icdf, _logcdf, _logprob
+from pymc.logprob.abstract import MeasurableOp, _icdf, _logcdf, _logprob
 from pymc.logprob.basic import logp
 from pymc.logprob.rewriting import logprob_rewrites_db
 from pymc.printing import str_for_dist
@@ -228,7 +228,7 @@ def __get__(self, instance, type_):
         return descr_get(instance, type_)
 
 
-class SymbolicRandomVariable(OpFromGraph):
+class SymbolicRandomVariable(MeasurableOp, OpFromGraph):
     """Symbolic Random Variable
 
     This is a subclasse of `OpFromGraph` which is used to encapsulate the symbolic
@@ -624,10 +624,6 @@ def dist(
         return rv_out
 
 
-# Let PyMC know that the SymbolicRandomVariable has a logprob.
-MeasurableVariable.register(SymbolicRandomVariable)
-
-
 @node_rewriter([SymbolicRandomVariable])
 def inline_symbolic_random_variable(fgraph, node):
     """

pymc/logprob/abstract.py

@@ -35,6 +35,7 @@
 #   SOFTWARE.
 
 import abc
+import warnings
 
 from collections.abc import Sequence
 from functools import singledispatch
@@ -46,6 +47,17 @@
 from pytensor.tensor.random.op import RandomVariable
 
 
+def __getattr__(name):
+    if name == "MeasurableVariable":
+        warnings.warn(
+            f"{name} has been deprecated in favor of MeasurableOp. Importing will fail in a future release.",
+            FutureWarning,
+        )
+        return MeasurableOpMixin
+
+    raise AttributeError(f"module {__name__} has no attribute {name}")
+
+
 @singledispatch
 def _logprob(
     op: Op,
@@ -131,14 +143,21 @@ def _icdf_helper(rv, value, **kwargs):
     return rv_icdf
 
 
-class MeasurableVariable(abc.ABC):
-    """A variable that can be assigned a measure/log-probability"""
+class MeasurableOp(abc.ABC):
+    """An operation whose outputs can be assigned a measure/log-probability"""
+
 
+MeasurableOp.register(RandomVariable)
 
-MeasurableVariable.register(RandomVariable)
 
+class MeasurableOpMixin(MeasurableOp):
+    """MeasurableOp Mixin with a distinctive string representation"""
 
-class MeasurableElemwise(Elemwise):
+    def __str__(self):
+        return f"Measurable{super().__str__()}"
+
+
+class MeasurableElemwise(MeasurableOpMixin, Elemwise):
     """Base class for Measurable Elemwise variables"""
 
     valid_scalar_types: tuple[MetaType, ...] = ()
@@ -150,6 +169,3 @@ def __init__(self, scalar_op, *args, **kwargs):
                 f"Acceptable types are {self.valid_scalar_types}"
             )
         super().__init__(scalar_op, *args, **kwargs)
-
-
-MeasurableVariable.register(MeasurableElemwise)

pymc/logprob/basic.py

@@ -56,7 +56,7 @@
 from pytensor.tensor.variable import TensorVariable
 
 from pymc.logprob.abstract import (
-    MeasurableVariable,
+    MeasurableOp,
     _icdf_helper,
     _logcdf_helper,
     _logprob,
@@ -522,7 +522,7 @@ def conditional_logp(
     while q:
         node = q.popleft()
 
-        if not isinstance(node.op, MeasurableVariable):
+        if not isinstance(node.op, MeasurableOp):
             continue
 
         q_values = [replacements[q_rv] for q_rv in node.outputs if q_rv in updated_rv_values]

pymc/logprob/checks.py

@@ -42,18 +42,15 @@
 from pytensor.tensor import TensorVariable
 from pytensor.tensor.shape import SpecifyShape
 
-from pymc.logprob.abstract import MeasurableVariable, _logprob, _logprob_helper
+from pymc.logprob.abstract import MeasurableOp, MeasurableOpMixin, _logprob, _logprob_helper
 from pymc.logprob.rewriting import PreserveRVMappings, measurable_ir_rewrites_db
 from pymc.logprob.utils import replace_rvs_by_values
 
 
-class MeasurableSpecifyShape(SpecifyShape):
+class MeasurableSpecifyShape(MeasurableOpMixin, SpecifyShape):
     """A placeholder used to specify a log-likelihood for a specify-shape sub-graph."""
 
 
-MeasurableVariable.register(MeasurableSpecifyShape)
-
-
 @_logprob.register(MeasurableSpecifyShape)
 def logprob_specify_shape(op, values, inner_rv, *shapes, **kwargs):
     (value,) = values
@@ -80,7 +77,7 @@ def find_measurable_specify_shapes(fgraph, node) -> list[TensorVariable] | None:
 
     if not (
         base_rv.owner
-        and isinstance(base_rv.owner.op, MeasurableVariable)
+        and isinstance(base_rv.owner.op, MeasurableOp)
         and base_rv not in rv_map_feature.rv_values
     ):
         return None  # pragma: no cover
@@ -99,13 +96,10 @@ def find_measurable_specify_shapes(fgraph, node) -> list[TensorVariable] | None:
 )
 
 
-class MeasurableCheckAndRaise(CheckAndRaise):
+class MeasurableCheckAndRaise(MeasurableOpMixin, CheckAndRaise):
     """A placeholder used to specify a log-likelihood for an assert sub-graph."""
 
 
-MeasurableVariable.register(MeasurableCheckAndRaise)
-
-
 @_logprob.register(MeasurableCheckAndRaise)
 def logprob_check_and_raise(op, values, inner_rv, *assertions, **kwargs):
     (value,) = values

pymc/logprob/cumsum.py

@@ -41,17 +41,14 @@
 from pytensor.tensor import TensorVariable
 from pytensor.tensor.extra_ops import CumOp
 
-from pymc.logprob.abstract import MeasurableVariable, _logprob, _logprob_helper
+from pymc.logprob.abstract import MeasurableOpMixin, _logprob, _logprob_helper
 from pymc.logprob.rewriting import PreserveRVMappings, measurable_ir_rewrites_db
 
 
-class MeasurableCumsum(CumOp):
+class MeasurableCumsum(MeasurableOpMixin, CumOp):
     """A placeholder used to specify a log-likelihood for a cumsum sub-graph."""
 
 
-MeasurableVariable.register(MeasurableCumsum)
-
-
 @_logprob.register(MeasurableCumsum)
 def logprob_cumsum(op, values, base_rv, **kwargs):
     """Compute the log-likelihood graph for a `Cumsum`."""

pymc/logprob/mixture.py

@@ -67,7 +67,8 @@
 
 from pymc.logprob.abstract import (
     MeasurableElemwise,
-    MeasurableVariable,
+    MeasurableOp,
+    MeasurableOpMixin,
     _logprob,
     _logprob_helper,
 )
@@ -217,7 +218,7 @@ def rv_pull_down(x: TensorVariable) -> TensorVariable:
     return fgraph.outputs[0]
 
 
-class MixtureRV(Op):
+class MixtureRV(MeasurableOpMixin, Op):
     """A placeholder used to specify a log-likelihood for a mixture sub-graph."""
 
     __props__ = ("indices_end_idx", "out_dtype", "out_broadcastable")
@@ -235,9 +236,6 @@ def perform(self, node, inputs, outputs):
         raise NotImplementedError("This is a stand-in Op.")  # pragma: no cover
 
 
-MeasurableVariable.register(MixtureRV)
-
-
 def get_stack_mixture_vars(
     node: Apply,
 ) -> tuple[list[TensorVariable] | None, int | None]:
@@ -457,13 +455,10 @@ def logprob_switch_mixture(op, values, switch_cond, component_true, component_fa
 )
 
 
-class MeasurableIfElse(IfElse):
+class MeasurableIfElse(MeasurableOpMixin, IfElse):
     """Measurable subclass of IfElse operator."""
 
 
-MeasurableVariable.register(MeasurableIfElse)
-
-
 @node_rewriter([IfElse])
 def useless_ifelse_outputs(fgraph, node):
     """Remove outputs that are shared across the IfElse branches."""
@@ -512,7 +507,7 @@ def find_measurable_ifelse_mixture(fgraph, node):
     base_rvs = assume_measured_ir_outputs(valued_rvs, base_rvs)
     if len(base_rvs) != op.n_outs * 2:
         return None
-    if not all(var.owner and isinstance(var.owner.op, MeasurableVariable) for var in base_rvs):
+    if not all(var.owner and isinstance(var.owner.op, MeasurableOp) for var in base_rvs):
         return None
 
     return MeasurableIfElse(n_outs=op.n_outs).make_node(if_var, *base_rvs).outputs

pymc/logprob/order.py

@@ -48,7 +48,7 @@
 from pytensor.tensor.variable import TensorVariable
 
 from pymc.logprob.abstract import (
-    MeasurableVariable,
+    MeasurableOpMixin,
     _logcdf_helper,
     _logprob,
     _logprob_helper,
@@ -59,20 +59,14 @@
 from pymc.pytensorf import constant_fold
 
 
-class MeasurableMax(Max):
+class MeasurableMax(MeasurableOpMixin, Max):
     """A placeholder used to specify a log-likelihood for a max sub-graph."""
 
 
-MeasurableVariable.register(MeasurableMax)
-
-
-class MeasurableMaxDiscrete(Max):
+class MeasurableMaxDiscrete(MeasurableOpMixin, Max):
     """A placeholder used to specify a log-likelihood for sub-graphs of maxima of discrete variables"""
 
 
-MeasurableVariable.register(MeasurableMaxDiscrete)
-
-
 @node_rewriter([Max])
 def find_measurable_max(fgraph: FunctionGraph, node: Apply) -> list[TensorVariable] | None:
     rv_map_feature = getattr(fgraph, "preserve_rv_mappings", None)
@@ -162,21 +156,15 @@ def max_logprob_discrete(op, values, base_rv, **kwargs):
     return logprob
 
 
-class MeasurableMaxNeg(Max):
+class MeasurableMaxNeg(MeasurableOpMixin, Max):
     """A placeholder used to specify a log-likelihood for a max(neg(x)) sub-graph.
     This shows up in the graph of min, which is (neg(max(neg(x)))."""
 
 
-MeasurableVariable.register(MeasurableMaxNeg)
-
-
-class MeasurableDiscreteMaxNeg(Max):
+class MeasurableDiscreteMaxNeg(MeasurableOpMixin, Max):
     """A placeholder used to specify a log-likelihood for sub-graphs of negative maxima of discrete variables"""
 
 
-MeasurableVariable.register(MeasurableDiscreteMaxNeg)
-
-
 @node_rewriter(tracks=[Max])
 def find_measurable_max_neg(fgraph: FunctionGraph, node: Apply) -> list[TensorVariable] | None:
     rv_map_feature = getattr(fgraph, "preserve_rv_mappings", None)

pymc/logprob/rewriting.py

@@ -81,7 +81,7 @@
 )
 from pytensor.tensor.variable import TensorVariable
 
-from pymc.logprob.abstract import MeasurableVariable
+from pymc.logprob.abstract import MeasurableOp
 from pymc.logprob.utils import DiracDelta
 
 inc_subtensor_ops = (IncSubtensor, AdvancedIncSubtensor, AdvancedIncSubtensor1)
@@ -139,7 +139,7 @@ def apply(self, fgraph):
                     continue
                 # This is where we filter only those nodes we care about:
                 # Nodes that have variables that we want to measure and are not yet measurable
-                if isinstance(node.op, MeasurableVariable):
+                if isinstance(node.op, MeasurableOp):
                     continue
                 if not any(out in rv_map_feature.needs_measuring for out in node.outputs):
                     continue
@@ -155,7 +155,7 @@ def apply(self, fgraph):
                             node_rewriter, "__name__", ""
                         )
                     # If we converted to a MeasurableVariable we're done here!
-                    if node not in fgraph.apply_nodes or isinstance(node.op, MeasurableVariable):
+                    if node not in fgraph.apply_nodes or isinstance(node.op, MeasurableOp):
                         # go to next node
                         break
 
@@ -274,7 +274,7 @@ def request_measurable(self, vars: Sequence[Variable]) -> list[Variable]:
             # Input vars or valued vars can't be measured for derived expressions
             if not var.owner or var in self.rv_values:
                 continue
-            if isinstance(var.owner.op, MeasurableVariable):
+            if isinstance(var.owner.op, MeasurableOp):
                 measurable.append(var)
             else:
                 self.needs_measuring.add(var)

pymc/logprob/scan.py

@@ -54,7 +54,7 @@
 from pytensor.tensor.variable import TensorVariable
 from pytensor.updates import OrderedUpdates
 
-from pymc.logprob.abstract import MeasurableVariable, _logprob
+from pymc.logprob.abstract import MeasurableOp, MeasurableOpMixin, _logprob
 from pymc.logprob.basic import conditional_logp
 from pymc.logprob.rewriting import (
     PreserveRVMappings,
@@ -66,16 +66,13 @@
 from pymc.logprob.utils import replace_rvs_by_values
 
 
-class MeasurableScan(Scan):
+class MeasurableScan(MeasurableOpMixin, Scan):
     """A placeholder used to specify a log-likelihood for a scan sub-graph."""
 
     def __str__(self):
         return f"Measurable({super().__str__()})"
 
 
-MeasurableVariable.register(MeasurableScan)
-
-
 def convert_outer_out_to_in(
     input_scan_args: ScanArgs,
     outer_out_vars: Iterable[TensorVariable],
@@ -288,7 +285,7 @@ def get_random_outer_outputs(
         io_type = oo_info.name[(oo_info.name.index("_", 6) + 1) :]
         inner_out_type = f"inner_out_{io_type}"
         io_var = getattr(scan_args, inner_out_type)[oo_info.index]
-        if io_var.owner and isinstance(io_var.owner.op, MeasurableVariable):
+        if io_var.owner and isinstance(io_var.owner.op, MeasurableOp):
             rv_vars.append((n, oo_var, io_var))
     return rv_vars