Difference between using one jit call and multiple jit calls.

[uduse]

Say I have two functions f and g that could be jit compiled. Say I have data x that always goes through f first then g. What are the differences between these two variants:

f = jit(f)
g = jit(g)
y = g(f(x))

versus

h = jit(lambda x: g(f(x)))
y = h(x)

Is there significant difference between these two variants and should I favor one variant over the other one in most cases?

[sharadmv]

In your first example, f and g will be compiled separately and there will be no compiler optimizations across them. In your second example, f and g will be part of one big compiled computation and there will be more opportunity for optimization. There will also only be one function call, which decreases dispatch overhead.

The caveat is that compiling two smaller functions can potentially be faster than compiling one big function, so you could save on compile time by splitting up your function into smaller compiled functions.

From personal experience, I recommend dropping a jit at the outermost level possible so the compiler has more opportunities to perform optimizations.