pymc-devs / nutpie

With pymc-devs/pymc#6736 merged, pymc rewrites the logp graph before applying the gradient. This removes some Ops from the graph, which should have some performance benefits in general but mainly target at pymc-devs/pymc#6717. Currently, nutpie calculates dlogp itself here:

nutpie/nutpie/compile_pymc.py

Line 197 in 69eb0ce

grads = pytensor.gradient.grad(logp, value_vars)

I think this can be replaced with a simple

grads = model.dlogp(value_vars)

Though I think we can also drop the value_vars, it only seems to be used for the grad.

Thanks for the suggestion :-)

I just included this in #42 (dc3cf2b)
It would probably be cleaner to reuse some of the compilation machinery in pymc instead of building this again...

Use pymc's dlogp with rewrites