Rewrite log1mexp(log1mexp(x)) to x

Author: maresb

pytensor/tensor/rewriting/math.py

@@ -625,6 +625,13 @@ def local_exp_log_nan_switch(fgraph, node):
         new_out = switch(ge(x, 0), log1p(-x), np.asarray(np.nan, old_out.dtype))
         return [new_out]
 
+    # Case for log1mexp(log1mexp(x)) -> x
+    if isinstance(prev_op, ps_math.Log1mexp) and isinstance(node_op, ps_math.Log1mexp):
+        x = x.owner.inputs[0]
+        old_out = node.outputs[0]
+        new_out = switch(le(x, 0), x, np.asarray(np.nan, old_out.dtype))
+        return [new_out]
+
 
 @register_canonicalize
 @register_specialize

tests/tensor/rewriting/test_math.py

@@ -2016,6 +2016,29 @@ def test_log1mexp_log(self):
         np.testing.assert_almost_equal(f(data_valid), expected)
         assert np.all(np.isnan(f(data_invalid)))
 
+    def test_log1mexp_log1mexp(self):
+        # log1mexp(log1mexp(x)) -> x
+        data_valid = -np.random.random((4, 3)).astype("float32")
+        data_valid[0, 0] = 0  # edge case
+        data_invalid = data_valid + 1.1
+
+        x = fmatrix()
+        f = function([x], log1mexp(log1mexp(x)), mode=self.mode.excluding("inplace"))
+        assert equal_computations(
+            f.maker.fgraph.outputs,
+            [
+                pt.switch(
+                    x <= np.array([[0]], dtype=np.int8),
+                    x,
+                    np.array([[np.nan]], dtype=np.float32),
+                )
+            ],
+        )
+
+        expected = data_valid
+        np.testing.assert_almost_equal(f(data_valid), expected)
+        assert np.all(np.isnan(f(data_invalid)))
+
     @pytest.mark.parametrize(
         ["nested_expression", "expected_switches"],
         [