Record assume_a as an assumption about the solved matrix - #2374
Open
jessegrabowski wants to merge 4 commits into
Open
Record assume_a as an assumption about the solved matrix#2374jessegrabowski wants to merge 4 commits into
jessegrabowski wants to merge 4 commits into
Conversation
The whole-graph pass at 0.8 only sees markers a user built, so it stays as the fast path; the local rewrite covers any a later rewrite introduces.
Member
|
Gosh these claude statements are really painful to read. Can you prompt it to use 2 sentences max or something? Propagate assumption from solve assume_a... nothing else was really said |
ricardoV94
reviewed
Aug 25, 2026
| } | ||
|
|
||
|
|
||
| def _normalize_assume_a(assume_a: str) -> str: |
Member
There was a problem hiding this comment.
I hate this, why lift stuff out of the Op? attributes and methods are a thing
| # _record_assume_a reads the dtype off ``a``, so it has to be a variable first. | ||
| a = pt.as_tensor_variable(a) | ||
| assume_a = _normalize_assume_a(assume_a) | ||
| a = _record_assume_a(a, assume_a) |
Member
Author
It's hard to get it down. I thought this was already a reasonable improvement from the default. |
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.
assume_apicked the LAPACK routine and nothing else, so a matrix the caller had already promised was positive definite or diagonal still looked opaque to every other op in the graph.solvenow records that promise as an assumption ona, which is the same thingassume()does. Diagonal gains the most, sincesolvelowers it to a division and the property was otherwise lost outright; a determinant, an inverse, a matmul, or a second solve against that matrix all drop from a dense op to an elementwise one.Rewrites build new solves too, so the recording can also happen mid-rewrite, after the single pass that resolves assumption markers has already run. Nothing drained those: the fact went nowhere, and the marker stayed in the graph where it blocks any rewrite trying to match across it. The drain is per-node work, so it now also runs as an ordinary local rewrite, with the one-shot pass kept as the fast path for the markers a user built. Two existing tests changed expectations because of that; they rewrite with
include=("canonicalize", ...), which never ran the one-shot pass, so markers had been leaking into graphs those tests called rewritten.