Derive logprob for Split operation #7875
Conversation
|
Classical mypy, feel free to review ignoring that, I'll fix it |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #7875 +/- ##
==========================================
+ Coverage 88.25% 92.94% +4.69%
==========================================
Files 116 116
Lines 18845 18875 +30
==========================================
+ Hits 16631 17544 +913
+ Misses 2214 1331 -883
🚀 New features to boost your workflow:
|
| # If the axis is over a dimension that was reduced in the logp (multivariate logp), | ||
| # We cannot split it into distinct entries. The mapping between values-densities breaks. | ||
| # We return the weighted logp by the split sizes. This is a good solution as any? | ||
| split_weights = splits / pt.sum(splits) |
There was a problem hiding this comment.
I think so? In MarginalMixture we decided to set the whole logp on the first entry, and zero for others, I like this approach more
| # axis=-2 (i.e., 0, - batch dimension) | ||
| x_parts = pt.split(x, splits_size=[2, 1], n_splits=2, axis=-2) | ||
| x_parts_vv = [x_part.clone() for x_part in x_parts] | ||
| logp_parts = list(conditional_logp(dict(zip(x_parts, x_parts_vv))).values()) |
There was a problem hiding this comment.
Do i understand this correctly that each part is conditioned on the values of all other parts?
Thinking about e.g. the MVN case, where if you split the vector and condition each split on the other, you get two new MVN distributions
There was a problem hiding this comment.
There's no marginalization going on, you can't evaluate the logp of only one part without providing the remaining ones. The only thing we do is join the value, get the logp, and split it again. We could argue that we don't want to do this for multivariate variables split along the core dimension, since there's no way you can split the logp (I did the weighing, but we can revert and raise NotImplemented)
📚 Documentation preview 📚: https://pymc--7875.org.readthedocs.build/en/7875/