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>
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. |
…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"). Both remain backward compatible (intermediates defaults to None). For a TransformedDistribution wrapped by to_event()/expand(), this lets log_prob reuse the sampled base value rather than inverting the transform, a potential speedup in certain cases. Include a test to make sure this actually works. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Eli Sennesh <elisennesh@astera.org>
Signed-off-by: Eli Sennesh <elisennesh@astera.org>
Got a test in there now for the if-yes branch, while the if-no branch is the default codepath exercised by most, though not quite all (hence the patch), existing code. |
|
Is there a way to convert the PR to pull from the |
|
|
||
| def fn(init_fn: Callable, *args, **kwargs) -> ArrayLike: | ||
| return init_fn(prng_key()) | ||
| def fn(*args, init_value: Callable = init_value, **kwargs) -> ArrayLike: |
There was a problem hiding this comment.
Seems like we dont need init_value here?
There was a problem hiding this comment.
It seems like we don't, but IME we actually do. Python behaves very nastily with closures and Jax behaves even worse. Generally IME the best way to actually close over something is like this: make the closed-over value the default for a pseudo-"argument", to make sure it becomes part of the function, gets passed by value, and survives closure conversion in the correct form.
If there was a way to demonstrate we don't need it, we could get rid of it, but again IME nasty bugs tend to show up after I mistake Python for Haskell in a case like this.
3 participants
init_valueto be something other than init-to-mean (75bebfc)init_valuefix (22aff3f)