From 061caeeb87584c567267796c50f4e2f9f8423780 Mon Sep 17 00:00:00 2001
From: jessegrabowski <jessegrabowski@gmail.com>
Date: Mon, 24 Aug 2026 20:30:38 -0500
Subject: [PATCH 1/8] Rewrite log(sqr(x)) to 2 * log(abs(x)) to avoid underflow

A squared product underflows to zero in float32 at a few hundred elements, sending log(abs(sqr(prod(x)))) to -inf; unwrapping the square lets the existing log(abs(prod(x))) -> sum(log(abs(x))) case fire before the product is ever materialized.
---
 pytensor/tensor/rewriting/math.py             | 51 ++++++++++++++++
 tests/tensor/rewriting/linalg/test_summary.py | 26 +++++++-
 tests/tensor/rewriting/test_math.py           | 59 +++++++++++++++++++
 3 files changed, 135 insertions(+), 1 deletion(-)

diff --git a/pytensor/tensor/rewriting/math.py b/pytensor/tensor/rewriting/math.py
index bb62d2a1ee..229f9a6e83 100644
--- a/pytensor/tensor/rewriting/math.py
+++ b/pytensor/tensor/rewriting/math.py
@@ -594,6 +594,27 @@ def local_sqrt_sqr(fgraph, node):
         return [new_out]
 
 
+@register_canonicalize
+@register_specialize
+@node_rewriter([pt_abs])
+def local_abs_sqr(fgraph, node):
+    [sqr_x] = node.inputs
+
+    if not (
+        sqr_x.owner
+        and isinstance(sqr_x.owner.op, Elemwise)
+        and isinstance(sqr_x.owner.op.scalar_op, ps.Sqr)
+    ):
+        return
+
+    # For complex x, abs(sqr(x)) is the squared modulus rather than the square
+    if sqr_x.dtype.startswith("complex"):
+        return
+
+    # Case for abs(sqr(x)) -> sqr(x)
+    return [sqr_x]
+
+
 @register_specialize
 @node_rewriter([log])
 def local_log_sqrt(fgraph, node):
@@ -617,6 +638,36 @@ def local_log_sqrt(fgraph, node):
     return [new_out]
 
 
+@register_stabilize
+@register_specialize
+@node_rewriter([log])
+def local_log_sqr(fgraph, node):
+    x = node.inputs[0]
+
+    if (
+        not x.owner
+        or not isinstance(x.owner.op, Elemwise)
+        or not isinstance(x.owner.op.scalar_op, ps.Sqr)
+    ):
+        return
+
+    x = x.owner.inputs[0]
+
+    # For complex x, abs(x) is real and would silently drop the imaginary part
+    if x.dtype.startswith("complex"):
+        return
+
+    # 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.
+    old_out = node.outputs[0]
+    new_out = mul(as_tensor_variable(2.0, dtype=old_out.dtype), log(pt_abs(x)))
+    if new_out.dtype != old_out.dtype:
+        new_out = cast(new_out, old_out.dtype)
+
+    copy_stack_trace(node.out, new_out)
+    return [new_out]
+
+
 @register_specialize
 @node_rewriter([exp, expm1, log1pexp, log1mexp])
 def local_exp_log_nan_switch(fgraph, node):
diff --git a/tests/tensor/rewriting/linalg/test_summary.py b/tests/tensor/rewriting/linalg/test_summary.py
index 517e178ebc..c213357148 100644
--- a/tests/tensor/rewriting/linalg/test_summary.py
+++ b/tests/tensor/rewriting/linalg/test_summary.py
@@ -10,6 +10,7 @@
 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
 
@@ -239,7 +240,7 @@ 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(
@@ -247,6 +248,11 @@ def test_slogdet_specialization():
             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,6 +273,24 @@ def test_local_log_prod_to_sum_log(original_fn, expected_fn):
     assert_equal_computations([rewritten], [expected])
 
 
+@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)
+
+
 @pytest.mark.parametrize(
     "expected, pos_tag",
     [
diff --git a/tests/tensor/rewriting/test_math.py b/tests/tensor/rewriting/test_math.py
index 0a2b845052..05f1756da5 100644
--- a/tests/tensor/rewriting/test_math.py
+++ b/tests/tensor/rewriting/test_math.py
@@ -2153,6 +2153,65 @@ def test_log_sqrt() -> None:
     )
 
 
+def test_log_sqr() -> None:
+    x = pt.tensor("x", shape=(None, None))
+    out = log(sqr(x))
+
+    out = rewrite_graph(out, include=["specialize"])
+
+    assert utt.assert_equal_computations(
+        [out],
+        [mul(pt.as_tensor_variable([[2.0]], dtype=x.dtype), log(pt_abs(x)))],
+    )
+
+
+def test_abs_sqr() -> None:
+    x = pt.tensor("x", shape=(None, None))
+    out = pt_abs(sqr(x))
+
+    out = rewrite_graph(out, include=["canonicalize", "specialize"])
+
+    assert utt.assert_equal_computations([out], [sqr(x)])
+
+
+def test_log_sqr_integer_input():
+    x = ivector("x")
+    out = log(sqr(x))
+
+    rewritten = rewrite_graph(out, include=["canonicalize", "stabilize", "specialize"])
+    assert rewritten.type.dtype == out.type.dtype
+
+    fn = function([x], rewritten, mode=Mode("py", None))
+    x_test = np.array([2, 3, 4], dtype="int32")
+    np.testing.assert_allclose(fn(x_test), np.log(x_test.astype("float64") ** 2))
+
+
+def test_log_sqr_extreme_magnitudes():
+    # sqr() saturates to inf above ~1e154 and to zero below ~1e-162 in float64
+    x = pt.vector("x")
+    fn = function([x], log(sqr(x)), mode="FAST_RUN")
+
+    x_test = np.array([1e200, 1e-200, 3.0])
+    np.testing.assert_allclose(fn(x_test), 2 * np.log(np.abs(x_test)))
+
+
+@pytest.mark.parametrize(
+    "original_fn",
+    [
+        pytest.param(lambda x: log(sqr(x)), id="log_sqr"),
+        pytest.param(lambda x: pt_abs(sqr(x)), id="abs_sqr"),
+    ],
+)
+def test_sqr_rewrites_skip_complex(original_fn):
+    # abs() of a complex input is real, so both rewrites would change the dtype
+    x = pt.vector("x", dtype="complex128")
+    out = original_fn(x)
+
+    rewritten = rewrite_graph(out, include=["canonicalize", "stabilize", "specialize"])
+
+    assert utt.assert_equal_computations([rewritten], [out])
+
+
 class TestSqrSqrt:
     def setup_method(self):
         mode = get_default_mode()

From 0ec77c301c09c7e3dda878158441cd5c30993609 Mon Sep 17 00:00:00 2001
From: jessegrabowski <jessegrabowski@gmail.com>
Date: Tue, 25 Aug 2026 16:12:29 -0500
Subject: [PATCH 2/8] Add a non_negative flag to unary scalar ops

---
 pytensor/scalar/basic.py   | 11 +++++++++++
 pytensor/scalar/math.py    |  3 +++
 tests/scalar/test_basic.py | 23 +++++++++++++++++++++++
 3 files changed, 37 insertions(+)

diff --git a/pytensor/scalar/basic.py b/pytensor/scalar/basic.py
index e3c088ff81..a8f6a8ab64 100644
--- a/pytensor/scalar/basic.py
+++ b/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):
diff --git a/pytensor/scalar/math.py b/pytensor/scalar/math.py
index f13297e572..d36f91054e 100644
--- a/pytensor/scalar/math.py
+++ b/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
diff --git a/tests/scalar/test_basic.py b/tests/scalar/test_basic.py
index 2aca416263..cb1264687d 100644
--- a/tests/scalar/test_basic.py
+++ b/tests/scalar/test_basic.py
@@ -4,6 +4,8 @@
 import pytest
 
 import pytensor
+import pytensor.scalar.basic as ps
+import pytensor.scalar.math as ps_math
 import pytensor.tensor as pt
 from pytensor.compile.mode import Mode, get_default_mode
 from pytensor.graph.fg import FunctionGraph
@@ -59,6 +61,7 @@
     true_div,
 )
 from pytensor.tensor import tensor_from_scalar
+from pytensor.tensor.elemwise import Elemwise
 from pytensor.tensor.type import fscalar, imatrix, matrix
 from tests.link.test_link import make_function
 
@@ -558,3 +561,23 @@ def test_pow_negative_base_fractional_exponent(mode):
             f"Expected numpy float, got {type(result)}: {result}"
         )
     assert np.isnan(result), f"Expected nan, got {result}"
+
+
+def test_non_negative_scalar_ops():
+    """Every op declaring ``non_negative`` must return >= 0, sign of zero included."""
+    flagged = {
+        op
+        for mod in (ps, ps_math)
+        for op in vars(mod).values()
+        if isinstance(op, ps.UnaryScalarOp) and op.non_negative
+    }
+    assert flagged, "no op carries the flag, the test is not exercising anything"
+
+    x_test = np.array([-1e3, -2.5, -1.0, -0.0, 0.0, 1e-8, 1.0, 2.5, 1e3])
+    for op in flagged:
+        x = pt.vector("x")
+        fn = pytensor.function([x], Elemwise(op)(x), mode=Mode("py", None))
+        out = fn(x_test)
+
+        assert (out >= 0).all(), f"{op} returned a negative value"
+        assert not np.signbit(out).any(), f"{op} returned -0.0"

From 3c92f8b05201c6e474974259516c02e4105d941f Mon Sep 17 00:00:00 2001
From: jessegrabowski <jessegrabowski@gmail.com>
Date: Tue, 25 Aug 2026 16:13:07 -0500
Subject: [PATCH 3/8] Fold abs() of any non-negative scalar op

---
 pytensor/tensor/rewriting/math.py   | 32 ++++++++++++--------
 tests/tensor/rewriting/test_math.py | 47 ++++++++++++++++++++++++++---
 2 files changed, 62 insertions(+), 17 deletions(-)

diff --git a/pytensor/tensor/rewriting/math.py b/pytensor/tensor/rewriting/math.py
index 229f9a6e83..38889b5ad0 100644
--- a/pytensor/tensor/rewriting/math.py
+++ b/pytensor/tensor/rewriting/math.py
@@ -594,25 +594,31 @@ def local_sqrt_sqr(fgraph, node):
         return [new_out]
 
 
+def _is_non_negative(var) -> bool:
+    """``True`` when ``var`` is known to be ``>= 0`` from the op that produced it.
+
+    Signed integers wrap on overflow (``sqr(int8(12)) == -112``) and complex values are
+    unordered, so the flag on the scalar op only carries over for floating-point and
+    unsigned outputs.
+    """
+    op = var.owner_op
+    if not (isinstance(op, Elemwise) and isinstance(op.scalar_op, ps.UnaryScalarOp)):
+        return False
+
+    return op.scalar_op.non_negative and var.dtype.startswith(("float", "uint"))
+
+
 @register_canonicalize
 @register_specialize
 @node_rewriter([pt_abs])
-def local_abs_sqr(fgraph, node):
-    [sqr_x] = node.inputs
-
-    if not (
-        sqr_x.owner
-        and isinstance(sqr_x.owner.op, Elemwise)
-        and isinstance(sqr_x.owner.op.scalar_op, ps.Sqr)
-    ):
-        return
+def local_useless_abs(fgraph, node):
+    # Case for abs(x) -> x, when x is already non-negative
+    [x] = node.inputs
 
-    # For complex x, abs(sqr(x)) is the squared modulus rather than the square
-    if sqr_x.dtype.startswith("complex"):
+    if not _is_non_negative(x):
         return
 
-    # Case for abs(sqr(x)) -> sqr(x)
-    return [sqr_x]
+    return [x]
 
 
 @register_specialize
diff --git a/tests/tensor/rewriting/test_math.py b/tests/tensor/rewriting/test_math.py
index 05f1756da5..50279eb8e7 100644
--- a/tests/tensor/rewriting/test_math.py
+++ b/tests/tensor/rewriting/test_math.py
@@ -87,6 +87,7 @@
     reciprocal,
     sigmoid,
     sign,
+    sin,
     sinh,
     softplus,
     sqr,
@@ -117,6 +118,7 @@
     local_reduce_chain,
     local_reduce_join,
     local_sum_prod_of_mul_or_div,
+    local_useless_abs,
     mul_canonizer,
     parse_mul_tree,
     perform_sigm_times_exp,
@@ -2165,13 +2167,50 @@ def test_log_sqr() -> None:
     )
 
 
-def test_abs_sqr() -> None:
+@pytest.mark.parametrize(
+    "inner_fn", [sqr, pt_abs, exp, sigmoid], ids=["sqr", "abs", "exp", "sigmoid"]
+)
+def test_useless_abs(inner_fn):
     x = pt.tensor("x", shape=(None, None))
-    out = pt_abs(sqr(x))
+    result = RewriteTester(
+        [x], [pt_abs(inner_fn(x))], include=None, custom_rewrite=local_useless_abs
+    )
+
+    result.assert_graph(inner_fn(x))
+    result.assert_eval(np.array([[1.0, -2.0], [3.0, -4.0]]))
+
+
+@pytest.mark.parametrize("inner_fn", [sin, neg], ids=["sin", "neg"])
+def test_useless_abs_sign_indefinite(inner_fn):
+    x = pt.tensor("x", shape=(None, None))
+    out = pt_abs(inner_fn(x))
+    result = RewriteTester([x], [out], include=None, custom_rewrite=local_useless_abs)
 
-    out = rewrite_graph(out, include=["canonicalize", "specialize"])
+    result.assert_graph(out)
+    result.assert_eval(np.array([[1.0, -2.0], [3.0, -4.0]]))
+
+
+def test_useless_abs_keeps_signed_zero():
+    # sqrt(-0.0) is -0.0, which abs() normalizes to +0.0, so dropping the abs would flip
+    # the sign of any downstream division
+    x = pt.vector("x")
+    out = pt_abs(sqrt(x))
+    result = RewriteTester([x], [out], include=None, custom_rewrite=local_useless_abs)
+
+    result.assert_graph(out)
+    [rewr_out] = result.rewr_fn(np.array([-0.0]))
+    assert not np.signbit(rewr_out)
+
+
+def test_useless_abs_signed_integer_overflow():
+    # sqr() wraps on signed ints (sqr(int8(12)) == -112), so its output is not
+    # non-negative and the abs() has to stay
+    x = pt.vector("x", dtype="int8")
+    out = pt_abs(sqr(x))
+    result = RewriteTester([x], [out], include=None, custom_rewrite=local_useless_abs)
 
-    assert utt.assert_equal_computations([out], [sqr(x)])
+    result.assert_graph(out)
+    result.assert_eval(np.array([12, 16, 3], dtype="int8"))
 
 
 def test_log_sqr_integer_input():

From e152eb7cec4208a356bb4eefaf39096b7c93576d Mon Sep 17 00:00:00 2001
From: jessegrabowski <jessegrabowski@gmail.com>
Date: Tue, 25 Aug 2026 16:15:09 -0500
Subject: [PATCH 4/8] Merge the log(sqrt(x)) and log(sqr(x)) rewrites

Sharing the tail means sharing the constant's dtype, which is what stops the 0.5 factor from truncating to 0 on integer input.

Closes #2379
---
 pytensor/tensor/rewriting/math.py   | 56 +++++++--------------
 tests/tensor/rewriting/test_math.py | 78 +++++++++++++++++------------
 2 files changed, 63 insertions(+), 71 deletions(-)

diff --git a/pytensor/tensor/rewriting/math.py b/pytensor/tensor/rewriting/math.py
index 38889b5ad0..3b42250794 100644
--- a/pytensor/tensor/rewriting/math.py
+++ b/pytensor/tensor/rewriting/math.py
@@ -621,52 +621,30 @@ def local_useless_abs(fgraph, node):
     return [x]
 
 
-@register_specialize
-@node_rewriter([log])
-def local_log_sqrt(fgraph, node):
-    x = node.inputs[0]
-
-    if (
-        not x.owner
-        or not isinstance(x.owner.op, Elemwise)
-        or not isinstance(x.owner.op.scalar_op, ps.Sqrt)
-    ):
-        return
-
-    # Case for log(sqrt(x)) -> 0.5 * log(x)
-    x = x.owner.inputs[0]
-    old_out = node.outputs[0]
-    new_out = mul(as_tensor_variable(0.5, dtype=x.dtype), log(x))
-    if new_out.dtype != old_out.dtype:
-        new_out = cast(new_out, old_out.dtype)
-
-    copy_stack_trace(node.out, new_out)
-    return [new_out]
-
-
 @register_stabilize
 @register_specialize
 @node_rewriter([log])
-def local_log_sqr(fgraph, node):
-    x = node.inputs[0]
-
-    if (
-        not x.owner
-        or not isinstance(x.owner.op, Elemwise)
-        or not isinstance(x.owner.op.scalar_op, ps.Sqr)
-    ):
-        return
+def local_log_sqrt_sqr(fgraph, node):
+    [x] = node.inputs
 
-    x = x.owner.inputs[0]
+    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))
 
-    # For complex x, abs(x) is real and would silently drop the imaginary part
-    if x.dtype.startswith("complex"):
-        return
+        case _:
+            return
 
-    # 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.
     old_out = node.outputs[0]
-    new_out = mul(as_tensor_variable(2.0, dtype=old_out.dtype), log(pt_abs(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)
 
diff --git a/tests/tensor/rewriting/test_math.py b/tests/tensor/rewriting/test_math.py
index 50279eb8e7..e88ac7b1ed 100644
--- a/tests/tensor/rewriting/test_math.py
+++ b/tests/tensor/rewriting/test_math.py
@@ -112,6 +112,7 @@
     local_div_switch_sink,
     local_grad_log_erfc_neg,
     local_greedy_distributor,
+    local_log_sqrt_sqr,
     local_mul_canonizer,
     local_mul_switch_sink,
     local_neg_to_mul,
@@ -2143,29 +2144,37 @@ def test_exp_log_nested(self, nested_expression, expected_switches):
         assert len(ops_graph) == expected_switches
 
 
-def test_log_sqrt() -> None:
+def test_log_sqrt():
     x = pt.tensor("x", shape=(None, None))
-    out = log(sqrt(x))
+    result = RewriteTester(
+        [x], [log(sqrt(x))], include=None, custom_rewrite=local_log_sqrt_sqr
+    )
 
-    out = rewrite_graph(out, include=["specialize"])
+    result.assert_graph(0.5 * log(x))
+    result.assert_eval(np.array([[1.0, 2.0], [3.0, 4.0]]))
 
-    assert utt.assert_equal_computations(
-        [out],
-        [mul(pt.as_tensor_variable([[0.5]], dtype=x.dtype), log(x))],
-    )
 
+def test_log_sqrt_integer_input():
+    # A 0.5 factor typed from the integer input would truncate to 0
+    x = ivector("x")
+    result = RewriteTester(
+        [x], [log(sqrt(x))], include=None, custom_rewrite=local_log_sqrt_sqr
+    )
 
-def test_log_sqr() -> None:
-    x = pt.tensor("x", shape=(None, None))
-    out = log(sqr(x))
+    result.assert_graph(0.5 * log(x))
+    assert result.rewr_fg.outputs[0].type.dtype == result.orig_fg.outputs[0].type.dtype
+    result.assert_eval(np.array([2, 3, 4], dtype="int32"))
 
-    out = rewrite_graph(out, include=["specialize"])
 
-    assert utt.assert_equal_computations(
-        [out],
-        [mul(pt.as_tensor_variable([[2.0]], dtype=x.dtype), log(pt_abs(x)))],
+def test_log_sqr():
+    x = pt.tensor("x", shape=(None, None))
+    result = RewriteTester(
+        [x], [log(sqr(x))], include=None, custom_rewrite=local_log_sqrt_sqr
     )
 
+    result.assert_graph(2.0 * log(pt_abs(x)))
+    result.assert_eval(np.array([[1.0, -2.0], [3.0, -4.0]]))
+
 
 @pytest.mark.parametrize(
     "inner_fn", [sqr, pt_abs, exp, sigmoid], ids=["sqr", "abs", "exp", "sigmoid"]
@@ -2215,40 +2224,45 @@ def test_useless_abs_signed_integer_overflow():
 
 def test_log_sqr_integer_input():
     x = ivector("x")
-    out = log(sqr(x))
-
-    rewritten = rewrite_graph(out, include=["canonicalize", "stabilize", "specialize"])
-    assert rewritten.type.dtype == out.type.dtype
+    result = RewriteTester(
+        [x], [log(sqr(x))], include=None, custom_rewrite=local_log_sqrt_sqr
+    )
 
-    fn = function([x], rewritten, mode=Mode("py", None))
-    x_test = np.array([2, 3, 4], dtype="int32")
-    np.testing.assert_allclose(fn(x_test), np.log(x_test.astype("float64") ** 2))
+    result.assert_graph(2.0 * log(pt_abs(x)))
+    assert result.rewr_fg.outputs[0].type.dtype == result.orig_fg.outputs[0].type.dtype
+    result.assert_eval(np.array([2, 3, 4], dtype="int32"))
 
 
 def test_log_sqr_extreme_magnitudes():
-    # sqr() saturates to inf above ~1e154 and to zero below ~1e-162 in float64
+    # sqr() saturates to inf above ~1e154 and to zero below ~1e-162 in float64, so the
+    # rewritten graph is deliberately not equivalent to the original one here
     x = pt.vector("x")
-    fn = function([x], log(sqr(x)), mode="FAST_RUN")
+    result = RewriteTester(
+        [x], [log(sqr(x))], include=None, custom_rewrite=local_log_sqrt_sqr
+    )
 
     x_test = np.array([1e200, 1e-200, 3.0])
-    np.testing.assert_allclose(fn(x_test), 2 * np.log(np.abs(x_test)))
+    [orig_out] = result.orig_fn(x_test)
+    [rewr_out] = result.rewr_fn(x_test)
+    assert np.isinf(orig_out).sum() == 2
+    np.testing.assert_allclose(rewr_out, 2 * np.log(np.abs(x_test)))
 
 
 @pytest.mark.parametrize(
-    "original_fn",
+    "original_fn, rewrite",
     [
-        pytest.param(lambda x: log(sqr(x)), id="log_sqr"),
-        pytest.param(lambda x: pt_abs(sqr(x)), id="abs_sqr"),
+        pytest.param(lambda x: log(sqr(x)), local_log_sqrt_sqr, id="log_sqr"),
+        pytest.param(lambda x: pt_abs(sqr(x)), local_useless_abs, id="abs_sqr"),
     ],
 )
-def test_sqr_rewrites_skip_complex(original_fn):
-    # abs() of a complex input is real, so both rewrites would change the dtype
+def test_sqr_rewrites_skip_complex(original_fn, rewrite):
+    # abs() of a complex input is real, so both rewrites would change the output dtype
     x = pt.vector("x", dtype="complex128")
     out = original_fn(x)
+    result = RewriteTester([x], [out], include=None, custom_rewrite=rewrite)
 
-    rewritten = rewrite_graph(out, include=["canonicalize", "stabilize", "specialize"])
-
-    assert utt.assert_equal_computations([rewritten], [out])
+    result.assert_graph(out)
+    result.assert_eval(np.array([1 + 2j, 0.5 - 1j]))
 
 
 class TestSqrSqrt:

From e4369e670d45262ae413ad1df17600b189a5edf5 Mon Sep 17 00:00:00 2001
From: jessegrabowski <jessegrabowski@gmail.com>
Date: Tue, 25 Aug 2026 16:52:23 -0500
Subject: [PATCH 5/8] Drop tautological dtype assertions from rewrite tests

FunctionGraph.replace filters the replacement through the old variable's type, so a rewrite that changed dtype would be rejected outright and these could never fail.
---
 tests/tensor/rewriting/test_math.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/tests/tensor/rewriting/test_math.py b/tests/tensor/rewriting/test_math.py
index e88ac7b1ed..7db2fc5a0b 100644
--- a/tests/tensor/rewriting/test_math.py
+++ b/tests/tensor/rewriting/test_math.py
@@ -2162,7 +2162,6 @@ def test_log_sqrt_integer_input():
     )
 
     result.assert_graph(0.5 * log(x))
-    assert result.rewr_fg.outputs[0].type.dtype == result.orig_fg.outputs[0].type.dtype
     result.assert_eval(np.array([2, 3, 4], dtype="int32"))
 
 
@@ -2229,7 +2228,6 @@ def test_log_sqr_integer_input():
     )
 
     result.assert_graph(2.0 * log(pt_abs(x)))
-    assert result.rewr_fg.outputs[0].type.dtype == result.orig_fg.outputs[0].type.dtype
     result.assert_eval(np.array([2, 3, 4], dtype="int32"))
 
 

From f1254255b205f4a9acd66e9901e33a62b4f9e5d5 Mon Sep 17 00:00:00 2001
From: jessegrabowski <jessegrabowski@gmail.com>
Date: Tue, 25 Aug 2026 16:52:36 -0500
Subject: [PATCH 6/8] Remove the non_negative flag property test

The hazard it covered is still pinned by test_useless_abs_keeps_signed_zero, which fails if Sqrt is ever flagged.
---
 tests/scalar/test_basic.py | 23 -----------------------
 1 file changed, 23 deletions(-)

diff --git a/tests/scalar/test_basic.py b/tests/scalar/test_basic.py
index cb1264687d..2aca416263 100644
--- a/tests/scalar/test_basic.py
+++ b/tests/scalar/test_basic.py
@@ -4,8 +4,6 @@
 import pytest
 
 import pytensor
-import pytensor.scalar.basic as ps
-import pytensor.scalar.math as ps_math
 import pytensor.tensor as pt
 from pytensor.compile.mode import Mode, get_default_mode
 from pytensor.graph.fg import FunctionGraph
@@ -61,7 +59,6 @@
     true_div,
 )
 from pytensor.tensor import tensor_from_scalar
-from pytensor.tensor.elemwise import Elemwise
 from pytensor.tensor.type import fscalar, imatrix, matrix
 from tests.link.test_link import make_function
 
@@ -561,23 +558,3 @@ def test_pow_negative_base_fractional_exponent(mode):
             f"Expected numpy float, got {type(result)}: {result}"
         )
     assert np.isnan(result), f"Expected nan, got {result}"
-
-
-def test_non_negative_scalar_ops():
-    """Every op declaring ``non_negative`` must return >= 0, sign of zero included."""
-    flagged = {
-        op
-        for mod in (ps, ps_math)
-        for op in vars(mod).values()
-        if isinstance(op, ps.UnaryScalarOp) and op.non_negative
-    }
-    assert flagged, "no op carries the flag, the test is not exercising anything"
-
-    x_test = np.array([-1e3, -2.5, -1.0, -0.0, 0.0, 1e-8, 1.0, 2.5, 1e3])
-    for op in flagged:
-        x = pt.vector("x")
-        fn = pytensor.function([x], Elemwise(op)(x), mode=Mode("py", None))
-        out = fn(x_test)
-
-        assert (out >= 0).all(), f"{op} returned a negative value"
-        assert not np.signbit(out).any(), f"{op} returned -0.0"

From 28287b991ef4aac687e8214ed24303f42e64e943 Mon Sep 17 00:00:00 2001
From: jessegrabowski <jessegrabowski@gmail.com>
Date: Tue, 25 Aug 2026 18:08:26 -0500
Subject: [PATCH 7/8] Teach _is_provably_non_negative about the non_negative
 flag

The flag case is non-strict only, since abs and sqr attain zero, and is restricted to float and unsigned outputs because signed integers wrap on overflow.
---
 pytensor/tensor/rewriting/math.py   | 22 ++++++----------------
 pytensor/tensor/subtensor.py        | 13 ++++++++++++-
 tests/tensor/rewriting/test_math.py | 21 +++++++++++++++++++++
 3 files changed, 39 insertions(+), 17 deletions(-)

diff --git a/pytensor/tensor/rewriting/math.py b/pytensor/tensor/rewriting/math.py
index 3b42250794..6005412b1a 100644
--- a/pytensor/tensor/rewriting/math.py
+++ b/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,20 +598,6 @@ def local_sqrt_sqr(fgraph, node):
         return [new_out]
 
 
-def _is_non_negative(var) -> bool:
-    """``True`` when ``var`` is known to be ``>= 0`` from the op that produced it.
-
-    Signed integers wrap on overflow (``sqr(int8(12)) == -112``) and complex values are
-    unordered, so the flag on the scalar op only carries over for floating-point and
-    unsigned outputs.
-    """
-    op = var.owner_op
-    if not (isinstance(op, Elemwise) and isinstance(op.scalar_op, ps.UnaryScalarOp)):
-        return False
-
-    return op.scalar_op.non_negative and var.dtype.startswith(("float", "uint"))
-
-
 @register_canonicalize
 @register_specialize
 @node_rewriter([pt_abs])
@@ -615,7 +605,7 @@ def local_useless_abs(fgraph, node):
     # Case for abs(x) -> x, when x is already non-negative
     [x] = node.inputs
 
-    if not _is_non_negative(x):
+    if not _is_provably_non_negative(x):
         return
 
     return [x]
diff --git a/pytensor/tensor/subtensor.py b/pytensor/tensor/subtensor.py
index 62a9879399..31cb18d726 100644
--- a/pytensor/tensor/subtensor.py
+++ b/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):
diff --git a/tests/tensor/rewriting/test_math.py b/tests/tensor/rewriting/test_math.py
index 7db2fc5a0b..3e4767eb26 100644
--- a/tests/tensor/rewriting/test_math.py
+++ b/tests/tensor/rewriting/test_math.py
@@ -2188,6 +2188,27 @@ def test_useless_abs(inner_fn):
     result.assert_eval(np.array([[1.0, -2.0], [3.0, -4.0]]))
 
 
+def test_useless_abs_unsigned_dtype():
+    x = pt.vector("x", dtype="uint8")
+    result = RewriteTester(
+        [x], [pt_abs(x)], include=None, custom_rewrite=local_useless_abs
+    )
+
+    result.assert_graph(x)
+    result.assert_eval(np.array([200, 3], dtype="uint8"))
+
+
+def test_useless_abs_of_clipped_input():
+    # maximum() is proven non-negative by the shared predicate, not by an op flag
+    x = pt.vector("x")
+    result = RewriteTester(
+        [x], [pt_abs(pt.maximum(x, 0))], include=None, custom_rewrite=local_useless_abs
+    )
+
+    result.assert_graph(pt.maximum(x, 0))
+    result.assert_eval(np.array([-2.0, 0.0, 3.0]))
+
+
 @pytest.mark.parametrize("inner_fn", [sin, neg], ids=["sin", "neg"])
 def test_useless_abs_sign_indefinite(inner_fn):
     x = pt.tensor("x", shape=(None, None))

From 7b110fb3dc5e64c503cf2e104e9299a72aa14be9 Mon Sep 17 00:00:00 2001
From: jessegrabowski <jessegrabowski@gmail.com>
Date: Tue, 25 Aug 2026 18:46:01 -0500
Subject: [PATCH 8/8] Use the shared predicate in local_log_prod_to_sum_log

Nothing in pytensor or pymc ever wrote the tag.positive hook the removed branch read, and the predicate narrows signed-integer operands, where sqr can wrap negative.
---
 pytensor/tensor/rewriting/linalg/summary.py   | 11 +++-----
 tests/tensor/rewriting/linalg/test_summary.py | 27 ++++---------------
 2 files changed, 9 insertions(+), 29 deletions(-)

diff --git a/pytensor/tensor/rewriting/linalg/summary.py b/pytensor/tensor/rewriting/linalg/summary.py
index 2e41fb3d1b..ae8a4e422c 100644
--- a/pytensor/tensor/rewriting/linalg/summary.py
+++ b/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,13 +27,14 @@
     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):
@@ -41,11 +42,7 @@ def local_log_prod_to_sum_log(fgraph, node):
             # 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
diff --git a/tests/tensor/rewriting/linalg/test_summary.py b/tests/tensor/rewriting/linalg/test_summary.py
index c213357148..f54c9f5926 100644
--- a/tests/tensor/rewriting/linalg/test_summary.py
+++ b/tests/tensor/rewriting/linalg/test_summary.py
@@ -12,7 +12,7 @@
 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():
@@ -291,30 +291,13 @@ def test_log_abs_sqr_prod_no_underflow(n):
     assert_allclose(fn(x_test), expected, rtol=1e-4)
 
 
-@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):
+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))
-    if pos_tag:
-        x.tag.positive = True
-
     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(