Miscellaneous bugfixes#2215
Conversation
Signed-off-by: Eli Sennesh <elisennesh@gmail.com>
…check Signed-off-by: Eli Sennesh <elisennesh@gmail.com>
Bugfix/mutables leaves
|
Hi! Please sync with master as we fixed the CI failure #2218 :) |
…alue() Signed-off-by: Eli Sennesh <elisennesh@gmail.com>
…dependent TransformedDistribution.sample_with_intermediates caches the pre-transform base sample so that log_prob(value, intermediates) can reuse it instead of recomputing transform.inv(value). That optimization was silently lost when the transformed distribution was wrapped: - Independent inherited the base no-op sample_with_intermediates (returning an empty []), so the cached value never reached the trace, and its log_prob did not accept intermediates. - ExpandedDistribution.log_prob ignored the intermediates it was passed (it carried a "TODO: utilize intermediates"). Independent only reinterprets batch dims as event dims and does not reshape the sample, so it can forward the base distribution's intermediates unchanged; ExpandedDistribution already produces intermediates in sample_with_intermediates and now forwards them in log_prob. Both remain backward compatible (intermediates defaults to None). For a TransformedDistribution wrapped by to_event()/expand() (e.g. a smoothed discrete site), this lets log_prob reuse the sampled base value rather than inverting the transform, an 11x speedup on an adaptive-inverse transform whose inverse is an iterative solve. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Eli Sennesh <elisennesh@astera.org>
All checks now pass 👍 . Let me know how you'd like to conduct the remainder of the review process! |
|
Thanks @esennesh . Most of it makes a lot of sense. Still, I wanna ask @fehiepsi for a review on the specific change
As I could be missing something. Maybe @esennesh, it's worth adding tests that show how this is supposed to work for each of the two |
Hmmm... I thought I had something on one of my other branches for submission that was gonna exercise this, but since I'm having a hard time finding that, I very much should add a test for exercising both branches. |
2 participants
init_valueto be something other than init-to-mean (75bebfc)init_valuefix (22aff3f)