search

pymc-devs / pymc

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

CheckParameterValues not replaced in inner graphs #6697

Open ricardoV94 opened 1 year ago

ricardoV94 commented 1 year ago

We include these rewrites when compiling logp graphs, but this won't replace inner logp graphs (e.g, inside Scan or not-inlined OpFromGraph)

https://github.com/pymc-devs/pymc/blob/371472dcf1c44f99a9ad4e7317606404b2a47a0e/pymc/pytensorf.py#L1142-L1144

The example in #6696 suffers from this. Sampling is quite faster when the rewrite is manually added

aerubanov commented 1 year ago

@ricardoV94 Looks like we can check inner graph for nodes that we need to replace, and if any found - create new Scan or OpFromGraph node with changed fgraph and replace old one. In you opinion will this approach work?

ricardoV94 commented 1 year ago

Actually for Scan there might be a shortcut, as Scan has it's own mode object. We can simple replace the old scan by a clone whose mode includes the rewrite we want to apply.

Eventually we want to get rid of this inner mode thing, but hopefully by that time we will have made Scan simpler so that's easier to manipulate anyway.

ricardoV94 commented 1 year ago

But we can do either way. The priority is more with Scan