I am working on speeding up a model with JIT compilation, as outlined here. I noticed in my own use case that there were diminishing returns to the JIT speedup as the problem size grew. I repeated the experiment with the simple Gaussian model at the link above, changing the number of observations from 100 to 1000 to 10000 and so on. At low problem sizes, the 2x speedup in the docs is replicated. At high problem sizes, though, the JIT-compiled SVI is no faster than the normal SVI.
Intuitively, the opposite should occur - compilation overhead should be amortized at high problem sizes. Has anyone observed this before?