Add PDMats extension#3030
Open
oschulz wants to merge 2 commits into
Open
Conversation
Traced values captured by closures that are called (but not otherwise referenced) inside @trace while/for bodies became block arguments of the while regions without matching while-loop operands, failing MLIR verification with 'expect operands to be compatible with body block arguments'. Treat single-name function calls in the loop condition and body as external symbols, so such closures become loop-carried arguments and the traced values they capture become proper while-loop operands. Created by generative AI.
PDMats.wsumsq and PDMats.invwsumsq use scalar accumulation loops, which degrade to one op per element under tracing (and fail under allowscalar(false)). They underlie e.g. the logpdf of diagonal-covariance MvNormal distributions. Provide broadcasted reductions for traced arrays via a package extension. Created by generative AI.
Contributor
Author
|
@wsmoses is this acceptable, in principle? |
wsmoses
reviewed
Jul 9, 2026
| # closure captures of the generated body function (which would result in | ||
| # region block arguments without matching operands). Only plain names can | ||
| # refer to local callables: | ||
| for symbols_state in (cond_symbols, body_symbols) |
wsmoses
reviewed
Jul 9, 2026
| # region block arguments without matching operands). Only plain names can | ||
| # refer to local callables: | ||
| for symbols_state in (cond_symbols, body_symbols) | ||
| for fn in symbols_state.funccalls |
Member
There was a problem hiding this comment.
@oschulz the pdmats extension looks pretty easy to give a quick approve to, this part here requires more scrutiny
Collaborator
There was a problem hiding this comment.
please make another pr for that change if not strictly required here. Also I would maybe apply that only if !isdefined(__module__, only(fn.parts)) to limit the number of arguments to each closures
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
3 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.
PDMats' scalar accumulation loops don't trace, use broadcasted reductions instead.
Came across this on while making MGVI.jl Reactant-compatible (in flight, depends on this PR).
AI-generated (Claude Fable at max effort), human reviewed.