With the functional API of Keras, it would definitely make sense. In fact I do think that imperative model definition would be great to have at some point in the future. We'll see :)
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...