You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
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 pymc-devs/pymc#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: