Rewrite log(sqr(x)) to 2 * log(abs(x)) to avoid underflow by jessegrabowski · Pull Request #2380 · pymc-devs/pytensor

jessegrabowski · 2026-08-25T01:52:40Z

Taking the log of a squared product materializes the product first, and in float32 a product of a few hundred terms underflows to zero once squared, so log(abs(sqr(prod(x)))) returns -inf. Shows up in GP code with a few hundred inducing points, on every backend.

log(sqrt(x)) and log(sqr(x)) are now one rewrite, so the square peels off and the existing log(abs(prod(x))) -> sum(log(abs(x))) case finishes the job before anything is squared. Dropping the abs is driven by a new non_negative flag on unary scalar ops instead of a hardcoded list, so abs(exp(x)) folds too. Sqrt deliberately doesn't carry the flag: sqrt(-0.0) is -0.0, and dropping the abs there would flip the sign of a downstream division.

x = pt.vector("x", dtype="float32")
f = pytensor.function([x], pt.log(pt.abs(pt.sqr(pt.prod(x)))))
f(np.full(250, 0.7, dtype="float32"))  # -inf before, -178.3 after

Closes #2379. Sharing one rewrite means sharing the constant's dtype, which is what stops the 0.5 factor truncating to 0 on integer input.

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.

ricardoV94 · 2026-08-25T06:06:04Z

+@register_canonicalize
+@register_specialize
+@node_rewriter([pt_abs])
+def local_abs_sqr(fgraph, node):


Include any scalar op that is non negative? Also use the new syntax x.owner_op_inputs?

I added a non_negative flag on ScalarOp just for the occasion!

ricardoV94 · 2026-08-25T06:06:55Z

+@register_stabilize
+@register_specialize
+@node_rewriter([log])
+def local_log_sqr(fgraph, node):


extend the sqrt rewrite instead? it's the same concept?

Sure why not

ricardoV94 · 2026-08-25T06:07:50Z

+        pytest.param(lambda x: pt_abs(sqr(x)), id="abs_sqr"),
+    ],
+)
+def test_sqr_rewrites_skip_complex(original_fn):


tell your bot about RewriteTest or whatever it's called...

: )

I love when people talk about how "we're going to have a tool that learns over time" as if these things aren't totally blank slate every time you boot them up.

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 pymc-devs#2379

ricardoV94 · 2026-08-25T21:26:43Z

+
+    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")


Not really a scalar test... Also do we test monotonicity? seems like over-zealous a bit?

yeah it's dumb

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.

The hazard it covered is still pinned by test_useless_abs_keeps_signed_zero, which fails if Sqrt is ever flagged.

ricardoV94 · 2026-08-25T22:02:18Z

-def local_log_sqrt(fgraph, node):
-    x = node.inputs[0]
+@node_rewriter([pt_abs])
+def local_useless_abs(fgraph, node):


does this supersede some other pre-existing more narrow rewrite?

no, but local_log_prod_to_sum_log was hand-rolling the check for non-negative Ops. I switched it to use _is_provably_non_negative (and by extension the new tag)

ricardoV94 · 2026-08-25T22:03:34Z

        return [new_out]


+def _is_non_negative(var) -> bool:


we have a smarter is non negative somewhere else (used for subtensor lift rewrites maybe)? that we should repurpose/extend with this? It includes stuff like dtype or dimshuffle / ravel ... bla bla

We have _is_provably_non_negative but it doesn't have any logic about functions like exp/softplus/whatever. I could extend it but I thought that was overkill

I added a case to _is_provably_non_negative that uses the new non_negative flag, then deleted the duplicate is_negative helper

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.

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.

ricardoV94 · 2026-08-27T21:49:39Z

        scalar_op = op.scalar_op
        if not strict and isinstance(scalar_op, Cast):
            return _is_provably_positive(var.owner.inputs[0], strict)
+        if (


follow up issue to move this outside of rewriting/subtensor.py

done. #2404

ricardoV94

Tests seem a bit overkill since the last iteration, but ok...

When a rewrite correctly does not fire, assert_eval compiles the unchanged graph twice and compares it to itself, so the int8 test now pins the actual wrapped output instead.

jessegrabowski · 2026-08-27T23:02:00Z

I trimmed back the tests, will merge when CI passes again.

ricardoV94 reviewed Aug 25, 2026

View reviewed changes

jessegrabowski added 3 commits August 25, 2026 16:12

Add a non_negative flag to unary scalar ops

0ec77c3

Fold abs() of any non-negative scalar op

3c92f8b

Merge the log(sqrt(x)) and log(sqr(x)) rewrites

e152eb7

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 pymc-devs#2379

ricardoV94 reviewed Aug 25, 2026

View reviewed changes

Comment thread tests/tensor/rewriting/test_math.py Outdated

ricardoV94 reviewed Aug 25, 2026

View reviewed changes

jessegrabowski added 2 commits August 25, 2026 16:52

Drop tautological dtype assertions from rewrite tests

e4369e6

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.

Remove the non_negative flag property test

f125425

The hazard it covered is still pinned by test_useless_abs_keeps_signed_zero, which fails if Sqrt is ever flagged.

ricardoV94 reviewed Aug 25, 2026

View reviewed changes

jessegrabowski added 2 commits August 25, 2026 18:08

Teach _is_provably_non_negative about the non_negative flag

28287b9

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.

Use the shared predicate in local_log_prod_to_sum_log

7b110fb

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.

ricardoV94 reviewed Aug 27, 2026

View reviewed changes

ricardoV94 approved these changes Aug 27, 2026

View reviewed changes

jessegrabowski added 3 commits August 27, 2026 18:47

Cover the non_negative flag in the _is_provably_positive tests

05b5b08

Drop rewrite tests that re-test the shared non-negative predicate

8925d5c

Replace tautological eval assertions in the abs rewrite tests

ee15eb2

When a rewrite correctly does not fire, assert_eval compiles the unchanged graph twice and compares it to itself, so the int8 test now pins the actual wrapped output instead.

jessegrabowski merged commit 597cdb7 into pymc-devs:main Aug 28, 2026
67 checks passed

jessegrabowski mentioned this pull request Aug 28, 2026

Move the _is_provably_positive helpers out of tensor/subtensor.py #2404

Open

Uh oh!

Conversation

jessegrabowski commented Aug 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

ricardoV94 Aug 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ricardoV94 left a comment

Choose a reason for hiding this comment

Uh oh!

jessegrabowski commented Aug 27, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

jessegrabowski commented Aug 25, 2026 •

edited

Loading

ricardoV94 Aug 25, 2026 •

edited

Loading