Full JIT impact on performance

[bhchiang]

How important is it to JIT all inner functions? (is there a way to do this recursively)

For example, if I'm doing vmap(jit(trace))(inp), should I bother adding the jit decorator to all of the smaller functions that trace calls? Thanks.

[bilal2vec]

I'm pretty sure you only need to jit the outermost function. also, you don't need to run @jit if you're already using pmap

[bhchiang]

Can you elaborate on why both of those are true? I can see JIT fusing everything through functions; I'm using vmap not pmap

[skye]

Nested jits have no effect if you're already in a jit context, so as long as you're jitting the outermost function, there's no benefit (or harm) to jitting inner functions. If you don't jit the inner functions, the tracing machinery doesn't "see" the function boundaries at all, and the whole computation is traced and compiled. If you do jit the inner functions, the tracing machinery now sees the function boundaries, but the inner functions are inlined into the final computation, so the whole computation is still compiled as if the inner jits weren't there.

@bkkaggle is correct that pmap also jit-compiles, so it's unnecessary to combine jit and pmap. vmap doesn't cause any compilation though, so you're correct to use jit with vmap.

[bhchiang]

Thanks for explaining!