Record assume_a as an assumption about the solved matrix - #2374
Record assume_a as an assumption about the solved matrix#2374jessegrabowski wants to merge 3 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.
|
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 |
| } | ||
|
|
||
|
|
||
| def _normalize_assume_a(assume_a: str) -> str: |
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) |
It's hard to get it down. I thought this was already a reasonable improvement from the default. |
c7a12cf to
dfc1775
Compare
|
Addressed all the feedback except the PR body. That stays to irritate you |
| @register_stabilize | ||
| @register_specialize | ||
| @node_rewriter([SpecifyAssumptions]) | ||
| def drain_specify_assumptions_node(fgraph, node): |
There was a problem hiding this comment.
I'm not sure about this... now there's 4 ways to trigger the assumption feature? It's like if the mere presence of SpecifyShape forced the ShapeFeature.
Why can't it be learned lazily (only one that eager skips solve is diagonal).
Also this explanation:
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.
It's only true because solve was rewritten to add the assume eagerly. So it's fixing a self-imposed issue.
There was a problem hiding this comment.
Alternatively you could have a single rewrite (stabilize maybe) that checks for solve and does the solve_assume -> assumption on the spot, iff the AssumptionFeature is triggered?
There was a problem hiding this comment.
Orthogonally, did you miss the direct solve_triangular call?
2 participants
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.