Implement measurable transforms for division, subtraction and negation

aeppl/transforms.py

@@ -10,8 +10,9 @@
 from aesara.graph.fg import FunctionGraph
 from aesara.graph.op import Op
 from aesara.graph.rewriting.basic import GraphRewriter, in2out, node_rewriter
-from aesara.scalar import Add, Exp, Log, Mul
+from aesara.scalar import Add, Exp, Log, Mul, Neg, Reciprocal, Sub, TrueDiv
 from aesara.tensor.elemwise import Elemwise
+from aesara.tensor.exceptions import NotScalarConstantError
 from aesara.tensor.rewriting.basic import (
     register_specialize,
     register_stabilize,
@@ -217,7 +218,7 @@ def apply(self, fgraph: FunctionGraph):
 class MeasurableTransform(MeasurableElemwise):
     """A placeholder used to specify a log-likelihood for a transformed measurable variable"""
 
-    valid_scalar_types = (Exp, Log, Add, Mul)
+    valid_scalar_types = (Exp, Log, Add, Mul, Reciprocal)
 
     # Cannot use `transform` as name because it would clash with the property added by
     # the `TransformValuesRewrite`
@@ -255,6 +256,92 @@ def measurable_transform_logprob(op: MeasurableTransform, values, *inputs, **kwa
     return input_logprob + jacobian
 
 
+@node_rewriter([Elemwise])
+def measurable_div_to_reciprocal_product(fgraph, node):
+    """Convert divisions involving `MeasurableVariable`s to product with reciprocal."""
+    if isinstance(node.op.scalar_op, TrueDiv):
+        measurable_vars = [
+            var
+            for var in node.inputs
+            if (var.owner and isinstance(var.owner.op, MeasurableVariable))
+        ]
+        if not measurable_vars:
+            return None  # pragma: no cover
+
+        rv_map_feature: Optional[PreserveRVMappings] = getattr(
+            fgraph, "preserve_rv_mappings", None
+        )
+        if rv_map_feature is None:
+            return None  # pragma: no cover
+
+        # Only apply this rewrite if there is one unvalued MeasurableVariable involved
+        if all(
+            measurable_var in rv_map_feature.rv_values
+            for measurable_var in measurable_vars
+        ):
+            return None  # pragma: no cover
+
+        numerator, denominator = node.inputs
+
+        # Check if numerator is 1
+        try:
+            if at.get_scalar_constant_value(numerator) == 1:
+                return [at.reciprocal(denominator)]
+        except NotScalarConstantError:
+            pass
+        return [at.mul(numerator, at.reciprocal(denominator))]
+
+
+@node_rewriter([Elemwise])
+def measurable_neg_to_product(fgraph, node):
+    """Convert negation of `MeasurableVariable`s to product with `-1`."""
+    if isinstance(node.op.scalar_op, Neg):
+        inp = node.inputs[0]
+        if not (inp.owner and isinstance(inp.owner.op, MeasurableVariable)):
+            return None
+
+        rv_map_feature: Optional[PreserveRVMappings] = getattr(
+            fgraph, "preserve_rv_mappings", None
+        )
+        if rv_map_feature is None:
+            return None  # pragma: no cover
+
+        # Only apply this rewrite if the variable is unvalued
+        if inp in rv_map_feature.rv_values:
+            return None  # pragma: no cover
+
+        return [at.mul(inp, -1.0)]
+
+
+@node_rewriter([Elemwise])
+def measurable_sub_to_neg(fgraph, node):
+    """Convert subtraction involving `MeasurableVariable`s to addition with neg"""
+    if isinstance(node.op.scalar_op, Sub):
+        measurable_vars = [
+            var
+            for var in node.inputs
+            if (var.owner and isinstance(var.owner.op, MeasurableVariable))
+        ]
+        if not measurable_vars:
+            return None  # pragma: no cover
+
+        rv_map_feature: Optional[PreserveRVMappings] = getattr(
+            fgraph, "preserve_rv_mappings", None
+        )
+        if rv_map_feature is None:
+            return None  # pragma: no cover
+
+        # Only apply this rewrite if there is one unvalued MeasurableVariable involved
+        if all(
+            measurable_var in rv_map_feature.rv_values
+            for measurable_var in measurable_vars
+        ):
+            return None  # pragma: no cover
+
+        minuend, subtrahend = node.inputs
+        return [at.add(minuend, at.neg(subtrahend))]
+
+
 @node_rewriter([Elemwise])
 def find_measurable_transforms(
     fgraph: FunctionGraph, node: Node
@@ -319,6 +406,8 @@ def find_measurable_transforms(
         transform = ExpTransform()
     elif isinstance(scalar_op, Log):
         transform = LogTransform()
+    elif isinstance(scalar_op, Reciprocal):
+        transform = ReciprocalTransform()
     elif isinstance(scalar_op, Add):
         transform_inputs = (measurable_input, at.add(*other_inputs))
         transform = LocTransform(
@@ -341,6 +430,30 @@ def find_measurable_transforms(
     return [transform_out]
 
 
+measurable_ir_rewrites_db.register(
+    "measurable_div_to_reciprocal_product",
+    measurable_div_to_reciprocal_product,
+    -5,
+    "basic",
+    "transform",
+)
+
+measurable_ir_rewrites_db.register(
+    "measurable_neg_to_product",
+    measurable_neg_to_product,
+    -5,
+    "basic",
+    "transform",
+)
+
+measurable_ir_rewrites_db.register(
+    "measurable_sub_to_neg",
+    measurable_sub_to_neg,
+    -5,
+    "basic",
+    "transform",
+)
+
 measurable_ir_rewrites_db.register(
     "find_measurable_transforms",
     find_measurable_transforms,
@@ -413,6 +526,19 @@ def log_jac_det(self, value, *inputs):
         return -at.log(value)
 
 
+class ReciprocalTransform(RVTransform):
+    name = "reciprocal"
+
+    def forward(self, value, *inputs):
+        return at.reciprocal(value)
+
+    def backward(self, value, *inputs):
+        return at.reciprocal(value)
+
+    def log_jac_det(self, value, *inputs):
+        return -2 * at.log(value)
+
+
 class IntervalTransform(RVTransform):
     name = "interval"
 

tests/test_transforms.py

@@ -591,17 +591,20 @@ def test_log_transform_rv():
 
 
 @pytest.mark.parametrize(
-    "rv_size, loc_type",
+    "rv_size, loc_type, addition",
     [
-        (None, at.scalar),
-        (2, at.vector),
-        ((2, 1), at.col),
+        (None, at.scalar, True),
+        (2, at.vector, False),
+        ((2, 1), at.col, True),
     ],
 )
-def test_loc_transform_rv(rv_size, loc_type):
+def test_loc_transform_rv(rv_size, loc_type, addition):
 
     loc = loc_type("loc")
-    y_rv = loc + at.random.normal(0, 1, size=rv_size, name="base_rv")
+    if addition:
+        y_rv = loc + at.random.normal(0, 1, size=rv_size, name="base_rv")
+    else:
+        y_rv = at.random.normal(0, 1, size=rv_size, name="base_rv") - at.neg(loc)
     y_rv.name = "y"
     y_vv = y_rv.clone()
 
@@ -619,17 +622,22 @@ def test_loc_transform_rv(rv_size, loc_type):
 
 
 @pytest.mark.parametrize(
-    "rv_size, scale_type",
+    "rv_size, scale_type, product",
     [
-        (None, at.scalar),
-        (1, at.TensorType("floatX", (True,))),
-        ((2, 3), at.matrix),
+        (None, at.scalar, True),
+        (1, at.TensorType("floatX", (True,)), True),
+        ((2, 3), at.matrix, False),
     ],
 )
-def test_scale_transform_rv(rv_size, scale_type):
+def test_scale_transform_rv(rv_size, scale_type, product):
 
     scale = scale_type("scale")
-    y_rv = at.random.normal(0, 1, size=rv_size, name="base_rv") * scale
+    if product:
+        y_rv = at.random.normal(0, 1, size=rv_size, name="base_rv") * scale
+    else:
+        y_rv = at.random.normal(0, 1, size=rv_size, name="base_rv") / at.reciprocal(
+            scale
+        )
     y_rv.name = "y"
     y_vv = y_rv.clone()
 
@@ -709,3 +717,41 @@ def test_invalid_broadcasted_transform_rv_fails():
     logp = joint_logprob({y_rv: y_vv})
     logp.eval({y_vv: [0, 0, 0, 0], loc: [0, 0, 0, 0]})
     assert False, "Should have failed before"
+
+
+@pytest.mark.parametrize("numerator", (1.0, 2.0))
+def test_reciprocal_rv_transform(numerator):
+    shape = 3
+    scale = 5
+    x_rv = numerator / at.random.gamma(shape, scale)
+    x_rv.name = "x"
+
+    x_vv = x_rv.clone()
+    x_logp_fn = aesara.function([x_vv], joint_logprob({x_rv: x_vv}))
+
+    x_test_val = 1.5
+    assert np.isclose(
+        x_logp_fn(x_test_val),
+        sp.stats.invgamma(shape, scale=scale * numerator).logpdf(x_test_val),
+    )
+
+
+def test_negated_rv_transform():
+    x_rv = -at.random.halfnormal()
+    x_rv.name = "x"
+
+    x_vv = x_rv.clone()
+    x_logp_fn = aesara.function([x_vv], joint_logprob({x_rv: x_vv}))
+
+    assert np.isclose(x_logp_fn(-1.5), sp.stats.halfnorm.logpdf(1.5))
+
+
+def test_subtracted_rv_transform():
+    # Choose base RV that is assymetric around zero
+    x_rv = 5.0 - at.random.normal(1.0)
+    x_rv.name = "x"
+
+    x_vv = x_rv.clone()
+    x_logp_fn = aesara.function([x_vv], joint_logprob({x_rv: x_vv}))
+
+    assert np.isclose(x_logp_fn(7.3), sp.stats.norm.logpdf(5.0 - 7.3, 1.0))