PyMC4 on Autograd?

Super hypothetical question… what if we used Autograd as the computational backend for PyMC4?

Happy for the idea to be shot down. One (and possibly only) advantage I could think of off the top of my head is that we’d never have to be swept away by a package going out of fad; Autograd is architecturally simple though a teeny bit of an overhead to maintain.

Anyways, just thought I’d ask, and maybe this’ll also revive some interest in the PyMC4 backend discussion. It feels like there’s a time limit before PyMC3 has to sunset, and I’d hate for this wonderful package to go away.

FWIW, I’m working on wrapping CuPy with autograd, so I’m thinking it might help with taking advantage of the GPU too!

I have never used autograd, but I heard they are a bit slow?

It is, but I think it’s mostly because autograd currently only auto diffs numpy code, which makes the matrix ops slower than it could be. (Mostly speaking from experience here; I built a keras-like thing using autograd when I was still in grad school, while I was still learning deep learning basics.)