add rewrite for log(sqrt(x))

Author: williambdean

pytensor/tensor/rewriting/math.py

@@ -552,6 +552,27 @@ def local_sqrt_sqr(fgraph, node):
         return [new_out]
 
 
+@register_canonicalize
+@node_rewriter([log])
+def local_log_sqrt(fgraph, node):
+    x = node.inputs[0]
+
+    if not (x.owner and isinstance(x.owner.op, Elemwise)):
+        return
+
+    prev_op = x.owner.op.scalar_op
+    node_op = node.op.scalar_op
+
+    if isinstance(prev_op, ps.Sqrt) and isinstance(node_op, ps.Log):
+        # Case for log(sqrt(x)) -> 0.5 * log(x)
+        x = x.owner.inputs[0]
+        old_out = node.outputs[0]
+        new_out = mul(0.5, log(x))
+        if new_out.dtype != old_out.dtype:
+            new_out = cast(new_out, old_out.dtype)
+        return [new_out]
+
+
 @register_specialize
 @node_rewriter([exp, expm1])
 def local_exp_log_nan_switch(fgraph, node):

tests/tensor/rewriting/test_math.py

@@ -1989,6 +1989,18 @@ def test_exp_log_nested(self, nested_expression, expected_switches):
         assert len(ops_graph) == expected_switches
 
 
+def test_log_sqrt() -> None:
+    x = pt.tensor("x", shape=(None, None))
+    out = log(sqrt(x))
+
+    out = rewrite_graph(out, include=["canonicalize", "specialize", "stabilize"])
+
+    assert equal_computations(
+        [out],
+        [mul(np.array([[0.5]]), log(x))],
+    )
+
+
 class TestSqrSqrt:
     def setup_method(self):
         mode = get_default_mode()