search

pymc-devs / pymc

Bayesian Modeling and Probabilistic Programming in Python
https://docs.pymc.io/
Other
8.47k stars 1.97k forks source link

Test `freeze_rv_and_dims` with `PartiallyObservedRV` #7316

Closed ricardoV94 closed 1 week ago

ricardoV94 commented 1 month ago

To use the JAX backend directly it's required that the boolean mask + advanced set subtensor in the graph of PartiallyObservedRVs get constant folded.

When using default mutable data/coords, this will not be the case, and the recommended solution is for users to make use of freeze_rv_and_dims. We should add a test that this actually works as suggested by @jessegrabowski

ricardoV94 commented 1 week ago

@jessegrabowski you have been vindicated. There was indeed a problem when combining the two things: #7388