I'm intrigued!... The kernel calling overhead and lack of any GPU while/scan/map/etc for Pytorch seems like a limitation, but I guess on 2nd thoughts you can still do all the keras fit/predict stuff and auto-connecting up the layers.
With TF's XLA compiler, they are slowly getting towards kernel fusion, which will then reduce launch overheads.
We have similar things in the works for pytorch: to quickly JIT at runtime the dynamic graph that is getting executed. More news on this will come when time-appropriate.
Also, have you looked at Numba to do the jitting? Probably best not to have yet another separately maintained python JIT.
https://discuss.pytorch.org/t/bayesian-computation-in-pytorc... https://discuss.pytorch.org/t/distribution-implementations/4...