Is there any reason to use nested jit?

[jatentaki]

In theory, @jit on the top-level function exposes the entire computation to XLA, enabling full optimization. In practice, if the top level function repeatedly uses some lower-level functions, annotating those could hint at the structure in the code, which the compiler will otherwise have to infer from the raw computation graph - this could lead to faster compilation. Is that the case with how jit works, or will the nested jit calls be "flattened" and only applied finally, at the top level (my original impression after reading the docs)? Similarly, is there any reason to prefer vmap(jit(f)) over jit(vmap(f))?

EDIT: I went with sprinkling @jit on some commonly reused subroutines and it changed compile time from so-long-I-never-got-any-results to about 1 minute (still long, but at least works) :)

[hawkinsp]

Right now, nested jit calls will be preserved as function calls in the IR that JAX generates, but will be flattened by XLA. In the future, this may not be true any more! At some point, XLA (or another compiler) may inline less aggressively.

[jatentaki]

To make sure I understood you: jit on the lower-level calls may improve compile performance (even if just by reducing the number of JAX IR generated and passed on to XLA) and cannot hurt because XLA will reason globally anyway?

[hawkinsp]

Yes, pretty much. Right now the compile time improvements will be small, but I would imagine that may change in the future.

[cyprienc]

Hi - sorry for unearthing this - is this answer still true? And, dumb question but, would the changes in the future increase or decrease the compilation time when decorating lower-level functions?
Thanks.

[francois-rozet]

For anyone coming back to this, jit on common low-level functions (e.g. neural network layers such as convolutions) is a great way to reduce compilation times for high-level functions (e.g. composition of layers).

For example,

>>> f = lambda x: sum(i * x for i in range(256))  # common function
>>> jf = jax.jit(f)
>>> %timeit jax.jit(lambda x: sum(f(x) for _ in range(16)))(3.14)
2.79 s ± 133 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
>>> %timeit jax.jit(lambda x: sum(jf(x) for _ in range(16)))(3.14)
36.5 ms ± 9.6 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)