pytensor
pytensor copied to clipboard
Add option to disable OpFromGraph.L_Op/R_Op from being wrapped in their own OpFromGraph
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