3 ms·
This is very interesting. The claim does sound too good to be true. Am I understanding you correctly that you are saying including the vmap operation in the tim
by geysersam 3y ago
This is very interesting. The claim does sound too good to be true. Am I understanding you correctly that you are saying including the vmap operation in the timing is wrong because it involves compilation time that could have been amortized over all the runs, and that the compilation time is considerable compared to the ode-solve itself?
- yablak 3y agoThere are two things going on. First, you're right that the vmap should have been done once outside the timing. But equally important, vmap(jit(...)) speed will generally be lower than jit(vmap(...)). There are many reasons for this. First, the former will loop iterations on the GPU in a serial fashion. Second, the internal jit makes optimization options opaque to Jax. For example, if there's a loop of matmuls inside main, that loop can be converted to a loop of einsums if you vmap first. It can also be fused into sometimes into a bigger operation that doesn't jump control variables back and forth between CPU and GPU between time steps. Between the two you both increase throughput and decrease latency. I think in Jax, jit(vmap(jit(...))) will also reoptimize the same way as jit(vmap(...)) but I'm not 100% certain.
- ubj 3y agoOn your last point, as long as you jit the topmost level, it doesn't matter whether or not you have inner jitted functions. The end result should be the same. Source: https://github.com/google/jax/discussions/5199#discussioncomment-219142 https://github.com/google/jax/discussions/5199#discussioncom...
- ChrisRackauckas 3y agoConfirmed in this collab notebook it doesn't make a tangible difference: https://colab.research.google.com/drive/1d7G-O5JX31lHbg7jTzzozbo5-Gp7DBEv#scrollTo=ExZj1Ak9UTuj https://colab.research.google.com/drive/1d7G-O5JX31lHbg7jTzz... .