2 ms·
On 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.
by ubj 3y ago
On 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... .