Add stabilization rewrite for log_diff_exp

Co-authored-by: Ricardo Vieira <[email protected]>

.rewrite

pytensor/tensor/rewriting/math.py

@@ -3604,6 +3604,14 @@ def local_reciprocal_1_plus_exp(fgraph, node):
 )
 register_stabilize(log1pmexp_to_log1mexp, name="log1pmexp_to_log1mexp")
 
+# log(exp(a) - exp(b)) -> a + log1mexp(b - a)
+logdiffexp_to_log1mexpdiff = PatternNodeRewriter(
+    (log, (sub, (exp, "x"), (exp, "y"))),
+    (add, "x", (log1mexp, (sub, "y", "x"))),
+    allow_multiple_clients=True,
+)
+register_stabilize(logdiffexp_to_log1mexpdiff, name="logdiffexp_to_log1mexpdiff")
+
 
 # log(sigmoid(x) / (1 - sigmoid(x))) -> x
 # i.e logit(sigmoid(x)) -> x

tests/tensor/rewriting/test_math.py

@@ -4136,3 +4136,42 @@ def test_log1mexp_stabilization():
         f(np.array([-0.8, -0.6], dtype=config.floatX)),
         np.log(1 - np.exp([-0.8, -0.6])),
     )
+
+
+def test_logdiffexp():
+    rng = np.random.default_rng(3559)
+    mode = Mode("py").including("stabilize").excluding("fusion")
+
+    x = fmatrix("x")
+    y = fmatrix("y")
+    f = function([x, y], log(exp(x) - exp(y)), mode=mode)
+
+    graph = f.maker.fgraph.toposort()
+    assert (
+        len(
+            [
+                node
+                for node in graph
+                if isinstance(node.op, Elemwise)
+                and isinstance(node.op.scalar_op, (aes.Exp, aes.Log))
+            ]
+        )
+        == 0
+    )
+    assert (
+        len(
+            [
+                node
+                for node in graph
+                if isinstance(node.op, Elemwise)
+                and isinstance(node.op.scalar_op, aes.Log1mexp)
+            ]
+        )
+        == 1
+    )
+
+    y_test = rng.normal(size=(3, 2)).astype("float32")
+    x_test = rng.normal(size=(3, 2)).astype("float32") + y_test.max()
+    np.testing.assert_almost_equal(
+        f(x_test, y_test), np.log(np.exp(x_test) - np.exp(y_test))
+    )