Rewrite log(sqr(x)) to 2 * log(abs(x)) to avoid underflow - #2380
Open
jessegrabowski wants to merge 1 commit into
Open
Rewrite log(sqr(x)) to 2 * log(abs(x)) to avoid underflow#2380jessegrabowski wants to merge 1 commit into
jessegrabowski wants to merge 1 commit into
Conversation
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
reviewed
Aug 25, 2026
| @register_canonicalize | ||
| @register_specialize | ||
| @node_rewriter([pt_abs]) | ||
| def local_abs_sqr(fgraph, node): |
Member
There was a problem hiding this comment.
Include any scalar op that is non negative? Also use the new syntax x.owner_op_inputs?
| @register_stabilize | ||
| @register_specialize | ||
| @node_rewriter([log]) | ||
| def local_log_sqr(fgraph, node): |
Member
There was a problem hiding this comment.
extend the sqrt rewrite instead? it's the same concept?
ricardoV94
reviewed
Aug 25, 2026
| pytest.param(lambda x: pt_abs(sqr(x)), id="abs_sqr"), | ||
| ], | ||
| ) | ||
| def test_sqr_rewrites_skip_complex(original_fn): |
Member
There was a problem hiding this comment.
tell your bot about RewriteTest or whatever it's called...
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.Two new rewrites,
local_log_sqrandlocal_abs_sqr, peel the square off so the existinglog(abs(prod(x))) -> sum(log(abs(x)))case can fire before anything is squared. Both skip complex inputs, whereabsis real-valued and would change the output dtype.