Derive logprob for exp2, log2, log10, log1p, expm1, log1mexp, log1pexp (softplus), and sigmoid transformations (#6826)

Author: LukeLB

pymc/logprob/transforms.py

@@ -61,10 +61,18 @@
     Erfc,
     Erfcx,
     Exp,
+    Exp2,
+    Expm1,
     Log,
+    Log1mexp,
+    Log1p,
+    Log2,
+    Log10,
     Mul,
     Pow,
+    Sigmoid,
     Sinh,
+    Softplus,
     Sqr,
     Sqrt,
     Tanh,
@@ -82,12 +90,20 @@
     erfc,
     erfcx,
     exp,
+    exp2,
+    expm1,
     log,
+    log1mexp,
+    log1p,
+    log2,
+    log10,
     mul,
     neg,
     pow,
     reciprocal,
+    sigmoid,
     sinh,
+    softplus,
     sqr,
     sqrt,
     sub,
@@ -569,8 +585,53 @@ def measurable_sub_to_neg(fgraph, node):
     return [pt.add(minuend, pt.neg(subtrahend))]
 
 
+@node_rewriter([log1p, softplus, log1mexp, log2, log10])
+def measurable_special_log_to_log(fgraph, node):
+    """Convert log1p, log1mexp, softplus, log2, log10 of `MeasurableVariable`s to log form."""
+    [inp] = node.inputs
+
+    if isinstance(node.op.scalar_op, Log1p):
+        return [pt.log(1 + inp)]
+    if isinstance(node.op.scalar_op, Softplus):
+        return [pt.log(1 + pt.exp(inp))]
+    if isinstance(node.op.scalar_op, Log1mexp):
+        return [pt.log(1 - pt.exp(inp))]
+    if isinstance(node.op.scalar_op, Log2):
+        return [pt.log(inp) / pt.log(2)]
+    if isinstance(node.op.scalar_op, Log10):
+        return [pt.log(inp) / pt.log(10)]
+
+
+@node_rewriter([expm1, sigmoid, exp2])
+def measurable_special_exp_to_exp(fgraph, node):
+    """Convert expm1, sigmoid, and exp2 of `MeasurableVariable`s to xp form."""
+    [inp] = node.inputs
+    if isinstance(node.op.scalar_op, Exp2):
+        return [pt.exp(pt.log(2) * inp)]
+    if isinstance(node.op.scalar_op, Expm1):
+        return [pt.add(pt.exp(inp), -1)]
+    if isinstance(node.op.scalar_op, Sigmoid):
+        return [1 / (1 + pt.exp(-inp))]
+
+
 @node_rewriter(
-    [exp, log, add, mul, pow, abs, sinh, cosh, tanh, arcsinh, arccosh, arctanh, erf, erfc, erfcx]
+    [
+        exp,
+        log,
+        add,
+        mul,
+        pow,
+        abs,
+        sinh,
+        cosh,
+        tanh,
+        arcsinh,
+        arccosh,
+        arctanh,
+        erf,
+        erfc,
+        erfcx,
+    ]
 )
 def find_measurable_transforms(fgraph: FunctionGraph, node: Node) -> Optional[List[Node]]:
     """Find measurable transformations from Elemwise operators."""
@@ -644,7 +705,6 @@ def find_measurable_transforms(fgraph: FunctionGraph, node: Node) -> Optional[Li
         transform = ScaleTransform(
             transform_args_fn=lambda *inputs: inputs[-1],
         )
-
     transform_op = MeasurableTransform(
         scalar_op=scalar_op,
         transform=transform,
@@ -692,6 +752,21 @@ def find_measurable_transforms(fgraph: FunctionGraph, node: Node) -> Optional[Li
     "transform",
 )
 
+measurable_ir_rewrites_db.register(
+    "measurable_special_log_to_log",
+    measurable_special_log_to_log,
+    "basic",
+    "transform",
+)
+
+measurable_ir_rewrites_db.register(
+    "measurable_special_exp_to_exp",
+    measurable_special_exp_to_exp,
+    "basic",
+    "transform",
+)
+
+
 measurable_ir_rewrites_db.register(
     "find_measurable_transforms",
     find_measurable_transforms,

tests/logprob/test_transforms.py

@@ -746,6 +746,8 @@ def test_chained_transform(self):
             ArcsinhTransform(),
             ArccoshTransform(),
             ArctanhTransform(),
+            LogTransform(),
+            ExpTransform(),
         ],
     )
     def test_check_jac_det(self, transform):
@@ -1105,7 +1107,6 @@ def test_cosh_rv_transform():
     # Something not centered around 0 is usually better
     base_rv = pt.random.normal(0.5, 1, size=(2,), name="base_rv")
     rv = pt.cosh(base_rv)
-
     vv = rv.clone()
     rv_logp = logp(rv, vv)
     with pytest.raises(NotImplementedError):
@@ -1118,14 +1119,46 @@ def test_cosh_rv_transform():
     expected_logp = pt.logaddexp(
         logp(base_rv, back_neg), logp(base_rv, back_pos)
     ) + transform.log_jac_det(vv)
-
     vv_test = np.array([0.25, 1.5])
     np.testing.assert_allclose(
         rv_logp.eval({vv: vv_test}),
         np.nan_to_num(expected_logp.eval({vv: vv_test}), nan=-np.inf),
     )
 
 
+TRANSFORMATIONS = {
+    "log1p": (pt.log1p, lambda x: pt.log(1 + x)),
+    "softplus": (pt.softplus, lambda x: pt.log(1 + pt.exp(x))),
+    "log1mexp": (pt.log1mexp, lambda x: pt.log(1 - pt.exp(x))),
+    "log2": (pt.log2, lambda x: pt.log(x) / pt.log(2)),
+    "log10": (pt.log10, lambda x: pt.log(x) / pt.log(10)),
+    "exp2": (pt.exp2, lambda x: pt.exp(pt.log(2) * x)),
+    "expm1": (pt.expm1, lambda x: pt.exp(x) - 1),
+    "sigmoid": (pt.sigmoid, lambda x: 1 / (1 + pt.exp(-x))),
+}
+
+
+@pytest.mark.parametrize("transform", TRANSFORMATIONS.keys())
+def test_special_log_exp_transforms(transform):
+    base_rv = pt.random.normal(name="base_rv")
+    vv = pt.scalar("vv")
+
+    transform_func, ref_func = TRANSFORMATIONS[transform]
+    transformed_rv = transform_func(base_rv)
+    ref_transformed_rv = ref_func(base_rv)
+
+    logp_test = logp(transformed_rv, vv)
+    logp_ref = logp(ref_transformed_rv, vv)
+
+    if transform in ["log2", "log10"]:
+        # in the cases of log2 and log10 floating point inprecision causes failure
+        # from equal_computations so evaluate logp and check all close instead
+        vv_test = np.array(0.25)
+        np.testing.assert_allclose(logp_ref.eval({vv: vv_test}), logp_test.eval({vv: vv_test}))
+    else:
+        assert equal_computations([logp_test], [logp_ref])
+
+
 @pytest.mark.parametrize("shift", [1.5, np.array([-0.5, 1, 0.3])])
 @pytest.mark.parametrize("scale", [2.0, np.array([1.5, 3.3, 1.0])])
 def test_multivariate_rv_transform(shift, scale):