Better handling of multiple output nodes in logprob inference - #8395
Draft
ricardoV94 wants to merge 3 commits into
Draft
Better handling of multiple output nodes in logprob inference#8395ricardoV94 wants to merge 3 commits into
ricardoV94 wants to merge 3 commits into
Conversation
ricardoV94
commented
Aug 9, 2026
ricardoV94
commented
Member
- Compatibility with pytensor 3.2
- Drop two stale pieces of logprob documentation
- Derive a joint density in a single call, however its values are reached
- Let a dims variable take part in a joint density
Documentation build overview
8 files changed ·
|
get_related_valued_nodes documents a `Returns` section listing an rv_node it does not return, and logprob_make_vector carries a TODO about a circular dependency that was resolved when the import it referred to moved to the module level; its twin in logprob_join was removed at the time and this one was left behind.
ricardoV94
force-pushed
the
measurable_multi_output
branch
from
August 10, 2026 07:16
bd38679 to
d14750d
Compare
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #8395 +/- ##
==========================================
- Coverage 91.85% 91.83% -0.02%
==========================================
Files 128 128
Lines 21259 21342 +83
==========================================
+ Hits 19527 19599 +72
- Misses 1732 1743 +11
🚀 New features to boost your workflow:
|
The logp helper bottomed out on one variable at a time, so a node whose density is joint over several outputs had it derived from a subset whenever one of the values arrived through a measurable chain -- exp, +1, a clip, anything. Such a target now leaves a DensityQuery behind, and a resolution pass answers every query about a node in one _logprob call. Answering one can only raise new queries about that node's ancestors, so a sinks-first pass resolves them all. Typing a query before its term exists needs to know which axes the density consumes, so MeasurableOp gains `supp_axes` -- the smallest piece of pymc-devs#6360 -- set by the rewrite that built the op; a query refuses rather than guess. Recursion into a measurable input now goes through request_logprob (renamed from _logprob_helper, which warns and forwards), which decides whether _logprob can be dispatched at all; it is dispatched directly only with every value of the node in hand. conditional_logp conditions by rewiring the graph it owns rather than substituting into clones, keeping node identity so that queries can be matched to their node.
The cast into xtensor space dispatched _logprob on the variable underneath rather
than going through the helper, so a dims value could never join the values of a
node whose density is joint over several of them.
Once it does, the term that comes back has to be labelled. Its dims were read off
positionally, on the assumption that a density is over the rightmost dims, so one
taken over a dim in the middle came back labelled by the dim it had just consumed
-- ("obs", "trial") reduced over "obs" was labelled ("obs",) rather than
("trial",). The cast op already carries the dims a density consumes as core_dims,
and the rewrite that builds it can now name them whenever the variable it wraps
says which axes its measure is over: they are counted from the right, and the op's
dims are laid out in the same order, so they index them directly. A variable that
says nothing leaves core_dims unset and keeps the positional guess, which is right
whenever the density really is over the rightmost dims.
ricardoV94
force-pushed
the
measurable_multi_output
branch
from
August 10, 2026 11:49
d14750d to
018c0a2
Compare
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
1 participant
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.