as core-devs, we dont plan to build-in something like Edward. However, folks in the community are brewing something:
https://discuss.pytorch.org/t/bayesian-computation-in-pytorc... https://discuss.pytorch.org/t/distribution-implementations/4...
Also, have you looked at Numba to do the jitting? Probably best not to have yet another separately maintained python JIT.