Cover more cases of log1mexp stabilization

pytensor/tensor/rewriting/math.py

@@ -3582,12 +3582,35 @@ def local_reciprocal_1_plus_exp(fgraph, node):
 register_specialize(local_1msigmoid)
 
 
-log1pmexp_to_log1mexp = PatternNodeRewriter(
-    (log1p, (neg, (exp, "x"))),
-    (log1mexp, "x"),
-    allow_multiple_clients=True,
-)
-register_stabilize(log1pmexp_to_log1mexp, name="log1pmexp_to_log1mexp")
+@register_stabilize
+@node_rewriter([log1p])
+def log1pmexp_to_log1mexp(fgraph, node):
+    """``log1p(-exp(x)) -> log1mexp(x)``
+    where "-" can be "neg" or any other expression detected by "is_neg"
+    """
+    (log1p_arg,) = node.inputs
+    exp_info = is_exp(log1p_arg)
+    if exp_info is not None:
+        exp_neg, exp_arg = exp_info
+        if exp_neg:
+            return [log1mexp(exp_arg)]
+        else:
+            return  # We could return [log1pexp(exp_arg)] here but that would conflict with log1pexp_to_softplus
+
+
+@register_stabilize
+@node_rewriter([log])
+def logmexpm1_to_log1mexp(fgraph, node):
+    """``log(-expm1(x)) -> log1mexp(x)``"""
+    (log_arg,) = node.inputs
+    if log_arg.owner and any(  #    ⬐ additional match to less-frequent expm1
+        i.owner and i.owner.op == expm1 for i in log_arg.owner.inputs
+    ):
+        neg_arg = is_neg(log_arg)
+        if neg_arg.owner and neg_arg.owner.op == expm1:
+            (expm1_arg,) = neg_arg.owner.inputs
+            return [log1mexp(expm1_arg)]
+
 
 # log(exp(a) - exp(b)) -> a + log1mexp(b - a)
 logdiffexp_to_log1mexpdiff = PatternNodeRewriter(

tests/tensor/rewriting/test_math.py

@@ -4438,11 +4438,22 @@ def test_local_add_neg_to_sub(first_negative):
     assert np.allclose(f(x_test, y_test), exp)
 
 
-def test_log1mexp_stabilization():
+@pytest.mark.parametrize(
+    "op_name",
+    ["log_1_minus_exp", "log1p_minus_exp", "log_minus_expm1", "log_minus_exp_minus_1"],
+)
+def test_log1mexp_stabilization(op_name):
     mode = Mode("py").including("stabilize")
 
     x = vector()
-    f = function([x], log(1 - exp(x)), mode=mode)
+    if op_name == "log_1_minus_exp":
+        f = function([x], log(1 - exp(x)), mode=mode)
+    elif op_name == "log1p_minus_exp":
+        f = function([x], log1p(-exp(x)), mode=mode)
+    elif op_name == "log_minus_expm1":
+        f = function([x], log(-expm1(x)), mode=mode)
+    elif op_name == "log_minus_exp_minus_1":
+        f = function([x], log(-(exp(x) - 1)), mode=mode)
 
     nodes = [node.op for node in f.maker.fgraph.toposort()]
     assert nodes == [pt.log1mexp]