pytensor icon indicating copy to clipboard operation
pytensor copied to clipboard

Add option to disable OpFromGraph.L_Op/R_Op from being wrapped in their own OpFromGraph

Open ricardoV94 opened this issue 10 months ago • 0 comments

The fact that we wrap a graph in OpFromGraph, doesn't mean we also want to wrap it's gradient. It can also lead to troubles as in https://github.com/pymc-devs/pymc/issues/7657

I suggest adding a wrap_grad_in_ofg: bool kwarg that toggles this behavior. The easiest way to get this change is to keep the behavior as is, but call the inline_ofg rewrite manually after calling it here:

https://github.com/pymc-devs/pytensor/blob/2f1d25a96e100380e87aebcf5f2ea7b678493ebd/pytensor/compile/builders.py#L612-L614

And here: https://github.com/pymc-devs/pytensor/blob/2f1d25a96e100380e87aebcf5f2ea7b678493ebd/pytensor/compile/builders.py#L689

ricardoV94 avatar Jan 25 '25 13:01 ricardoV94