Rewrite log(sqr(x)) to 2 * log(abs(x)) to avoid underflow

pytensor/scalar/basic.py

@@ -1296,6 +1296,11 @@ class UnaryScalarOp(ScalarOp):
     monotonic_increasing = False
     monotonic_decreasing = False
 
+    # Set only when the output is >= 0 for every real input, including the sign of zero
+    # (`sqrt` returns -0.0 for -0.0, so it does not qualify). Consumers must still check
+    # the dtype: signed integers wrap on overflow and complex values are unordered.
+    non_negative = False
+
     def c_code_contiguous(self, node, name, inputs, outputs, sub):
         (x,) = inputs
         (z,) = outputs
@@ -2523,6 +2528,7 @@ def cast(x, dtype):
 
 class Abs(UnaryScalarOp):
     preserves_zero = True
+    non_negative = True
     nfunc_spec = ("abs", 1, 1)
 
     def make_node(self, x):
@@ -3084,6 +3090,7 @@ def c_code(self, node, name, inputs, outputs, sub):
 
 class Exp(UnaryScalarOp):
     monotonic_increasing = True
+    non_negative = True
     nfunc_spec = ("exp", 1, 1)
     amd_float32 = "amd_vrsa_expf"
     amd_float64 = "amd_vrda_exp"
@@ -3123,6 +3130,7 @@ def c_code(self, node, name, inputs, outputs, sub):
 
 class Exp2(UnaryScalarOp):
     monotonic_increasing = True
+    non_negative = True
     nfunc_spec = ("exp2", 1, 1)
 
     def impl(self, x):
@@ -3201,6 +3209,7 @@ def c_code_cache_version(self):
 
 class Sqr(UnaryScalarOp):
     preserves_zero = True
+    non_negative = True
     nfunc_spec = ("square", 1, 1)
 
     def impl(self, x):
@@ -3546,6 +3555,8 @@ class Cosh(UnaryScalarOp):
 
     """
 
+    non_negative = True
+
     nfunc_spec = ("cosh", 1, 1)
 
     def impl(self, x):

pytensor/scalar/math.py

@@ -90,6 +90,7 @@ def c_code(self, node, name, inp, out, sub):
 
 class Erfc(UnaryScalarOp):
     monotonic_decreasing = True
+    non_negative = True
     nfunc_spec = ("scipy.special.erfc", 1, 1)
 
     def impl(self, x):
@@ -1170,6 +1171,7 @@ class Sigmoid(UnaryScalarOp):
     """
 
     monotonic_increasing = True
+    non_negative = True
     nfunc_spec = ("scipy.special.expit", 1, 1)
 
     def impl(self, x):
@@ -1225,6 +1227,7 @@ class Softplus(UnaryScalarOp):
     """
 
     monotonic_increasing = True
+    non_negative = True
 
     def impl(self, x):
         # If x is an int8 or uint8, numpy.exp will compute the result in

pytensor/tensor/rewriting/linalg/summary.py

@@ -10,7 +10,7 @@
     copy_stack_trace,
     node_rewriter,
 )
-from pytensor.scalar.basic import Abs, Exp, Log, Sign, Sqr
+from pytensor.scalar.basic import Abs, Log, Sign, Sqr
 from pytensor.tensor.basic import ones
 from pytensor.tensor.blockwise import Blockwise
 from pytensor.tensor.elemwise import Elemwise
@@ -27,25 +27,22 @@
     register_stabilize,
 )
 from pytensor.tensor.rewriting.linalg.utils import matrix_diagonal_product
+from pytensor.tensor.subtensor import _is_provably_non_negative
 
 
 @register_stabilize
 @register_specialize
 @node_rewriter([log])
 def local_log_prod_to_sum_log(fgraph, node):
-    """Rewrite log(prod(x)) as sum(log(x)), when x is known to be positive."""
+    """Rewrite log(prod(x)) as sum(log(x)), when x is known to be non-negative."""
     [p] = node.inputs
     match p.owner_op_and_inputs:
         case (Prod(axis=axis), x):
             # TODO: have a reduction like prod and sum that simply
             # returns the sign of the prod multiplication.
 
             # TODO: The product of diagonals of a Cholesky(A) are also strictly positive
-            match x.owner_op:
-                case Elemwise(Abs() | Sqr() | Exp()):
-                    return [log(x).sum(axis=axis)]
-
-            if getattr(x.tag, "positive", False):
+            if _is_provably_non_negative(x):
                 return [log(x).sum(axis=axis)]
 
         # Special case for log(abs(prod(x))) -> sum(log(abs(x))) that shows up in slogdet

pytensor/tensor/rewriting/math.py

@@ -110,7 +110,11 @@
 from pytensor.tensor.rewriting.blockwise import blockwise_of
 from pytensor.tensor.rewriting.elemwise import apply_local_dimshuffle_lift
 from pytensor.tensor.shape import Shape, Shape_i, specify_shape
-from pytensor.tensor.subtensor import Subtensor, _is_provably_positive
+from pytensor.tensor.subtensor import (
+    Subtensor,
+    _is_provably_non_negative,
+    _is_provably_positive,
+)
 from pytensor.tensor.type import (
     complex_dtypes,
     uint_dtypes,
@@ -594,22 +598,43 @@ def local_sqrt_sqr(fgraph, node):
         return [new_out]
 
 
+@register_canonicalize
 @register_specialize
-@node_rewriter([log])
-def local_log_sqrt(fgraph, node):
-    x = node.inputs[0]
+@node_rewriter([pt_abs])
+def local_useless_abs(fgraph, node):
+    # Case for abs(x) -> x, when x is already non-negative
+    [x] = node.inputs
 
-    if (
-        not x.owner
-        or not isinstance(x.owner.op, Elemwise)
-        or not isinstance(x.owner.op.scalar_op, ps.Sqrt)
-    ):
+    if not _is_provably_non_negative(x):
         return
 
-    # Case for log(sqrt(x)) -> 0.5 * log(x)
-    x = x.owner.inputs[0]
+    return [x]
+
+
+@register_stabilize
+@register_specialize
+@node_rewriter([log])
+def local_log_sqrt_sqr(fgraph, node):
+    [x] = node.inputs
+
+    match x.owner_op_and_inputs:
+        # Case for log(sqrt(x)) -> 0.5 * log(x)
+        case (Elemwise(ps.Sqrt()), inner):
+            factor, inner_out = 0.5, log(inner)
+
+        # Case for log(sqr(x)) -> 2 * log(abs(x)), which never materializes the square,
+        # so a value that would over- or underflow when squared survives the log.
+        case (Elemwise(ps.Sqr()), inner):
+            # abs() of a complex input is real and would drop the imaginary part
+            if inner.dtype.startswith("complex"):
+                return
+            factor, inner_out = 2.0, log(pt_abs(inner))
+
+        case _:
+            return
+
     old_out = node.outputs[0]
-    new_out = mul(as_tensor_variable(0.5, dtype=x.dtype), log(x))
+    new_out = mul(as_tensor_variable(factor, dtype=old_out.dtype), inner_out)
     if new_out.dtype != old_out.dtype:
         new_out = cast(new_out, old_out.dtype)
 

pytensor/tensor/subtensor.py

@@ -22,6 +22,7 @@
     ScalarMaximum,
     ScalarMinimum,
     ScalarVariable,
+    UnaryScalarOp,
 )
 from pytensor.tensor import (
     TensorLike,
@@ -253,12 +254,15 @@ def _is_provably_positive(var, strict: bool = True) -> bool:
     - ``minimum(a, b)`` when both ``a`` and ``b`` are positive.
     - ``maximum(a, b)`` when at least one of ``a``, ``b`` is positive.
 
-    Three further cases prove non-negativity but not strict positivity, so they
+    Four further cases prove non-negativity but not strict positivity, so they
     are recognized only when ``strict=False``:
 
     - Unsigned-integer dtype (a ``uint`` may be 0).
     - ``Shape`` / ``Shape_i`` outputs (a dimension may be 0).
     - ``Cast`` of a non-negative input (a float :math:`0 < x < 1` truncates to 0).
+    - A unary scalar op declaring ``non_negative`` (``abs``, ``sqr``, ``exp``, ...),
+      restricted to float and unsigned outputs: signed integers wrap on overflow
+      (``sqr(int8(12)) == -112``) and complex values are unordered.
 
     Parameters
     ----------
@@ -293,6 +297,13 @@ def _is_provably_positive(var, strict: bool = True) -> bool:
         scalar_op = op.scalar_op
         if not strict and isinstance(scalar_op, Cast):
             return _is_provably_positive(var.owner.inputs[0], strict)
+        if (
+            not strict
+            and isinstance(scalar_op, UnaryScalarOp)
+            and scalar_op.non_negative
+            and var.type.dtype.startswith(("float", "uint"))
+        ):
+            return True
         if isinstance(scalar_op, ScalarMinimum):
             return all(_is_provably_positive(i, strict) for i in var.owner.inputs)
         if isinstance(scalar_op, ScalarMaximum):

tests/tensor/rewriting/linalg/test_summary.py

@@ -10,8 +10,9 @@
 from pytensor.tensor.linalg.decomposition import lu, qr, svd
 from pytensor.tensor.linalg.decomposition.cholesky import cholesky
 from pytensor.tensor.linalg.summary import Det, SLogDet, det
+from pytensor.tensor.math import Prod
 from pytensor.tensor.type import matrix
-from tests.unittest_tools import assert_equal_computations
+from tests.unittest_tools import RewriteTester, assert_equal_computations
 
 
 def test_det_of_cholesky():
@@ -239,14 +240,19 @@ def test_slogdet_specialization():
         ),
         pytest.param(
             lambda x: pt.log(pt.prod(x**2)),
-            lambda x: pt.sum(pt.log(pt.sqr(x))),
+            lambda x: np.float64(2.0) * pt.sum(pt.log(pt.abs(x))),
             id="log_prod_sqr",
         ),
         pytest.param(
             lambda x: pt.log(pt.abs(pt.prod(x))),
             lambda x: pt.sum(pt.log(pt.abs(x))),
             id="log_abs_prod",
         ),
+        pytest.param(
+            lambda x: pt.log(pt.abs(pt.sqr(pt.prod(x)))),
+            lambda x: np.float64(2.0) * pt.sum(pt.log(pt.abs(x))),
+            id="log_abs_sqr_prod",
+        ),
         pytest.param(
             lambda x: pt.log(pt.prod(pt.abs(x), axis=0)),
             lambda x: pt.sum(pt.log(pt.abs(x)), axis=0),
@@ -267,30 +273,31 @@ def test_local_log_prod_to_sum_log(original_fn, expected_fn):
     assert_equal_computations([rewritten], [expected])
 
 
-@pytest.mark.parametrize(
-    "expected, pos_tag",
-    [
-        pytest.param(
-            lambda x: pt.sum(pt.log(x)),
-            True,
-            id="local_log_prod_to_sum_log_positive_tag",
-        ),
-        pytest.param(
-            lambda x: pt.log(pt.prod(x)),
-            False,
-            id="local_log_prod_to_sum_log_no_rewrite",
-        ),
-    ],
-)
-def test_local_log_prod_to_sum_log_positive_tag(expected, pos_tag):
-    x = pt.tensor("x", shape=(3, 4))
-    if pos_tag:
-        x.tag.positive = True
+@pytest.mark.parametrize("n", [200, 250, 300])
+def test_log_abs_sqr_prod_no_underflow(n):
+    """A float32 product of this many terms is zero once squared, so it must not be materialized."""
+    x = pt.vector("x", dtype="float32")
+    out = pt.log(pt.abs(pt.sqr(pt.prod(x))))
+
+    fn = function([x], out, mode="FAST_RUN")
+    assert not any(isinstance(node.op, Prod) for node in fn.maker.fgraph.apply_nodes)
+
+    rng = np.random.default_rng(sum(map(ord, "log_sqr_prod")))
+    # Mixed signs, so a rewrite that dropped the abs would return nan
+    x_test = (rng.uniform(0.5, 0.9, size=n) * rng.choice([-1, 1], size=n)).astype(
+        "float32"
+    )
+    expected = 2 * np.sum(np.log(np.abs(x_test.astype("float64"))))
+    assert_allclose(fn(x_test), expected, rtol=1e-4)
+
 
+def test_local_log_prod_to_sum_log_unknown_sign():
+    """A sign-indefinite operand leaves log(prod(x)) alone."""
+    x = pt.tensor("x", shape=(3, 4))
     out = pt.log(pt.prod(x))
+    result = RewriteTester([x], [out], include=["stabilize", "specialize"])
 
-    rewritten = rewrite_graph(out, include=["stabilize", "specialize"])
-    assert_equal_computations([rewritten], [expected(x)])
+    result.assert_graph(out)
 
 
 @pytest.mark.parametrize(