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)

[cjchristopher]

Revisiting this again in 2025, and I don't think this is the case anymore with the recent pulling of pjit.cc across into jaxlib.
Could someone confirm, and possibly adjust this answer? (also see my recent Issue here: #30270 )

Notably, if I jax.jit(jax.vmap(f)), where say

def f(in_vars):
    @jax.jit
    def sub_fun(vars):
         return jnp.some_op(vars)
   x = jnp.some_op(in_vars)
   y = sub_fun(x)
   return y

And turn on verbose logging (TF_CPP_MAX_VLOG_LEVEL=2 JAX_COMPILER_DETAILED_LOGGING_MIN_OPS=0 TF_CPP_MIN_LOG_LEVEL=0 JAX_TRACEBACK_FILTERING=off)

You end up with many (function name e.g. annotation mine):

I jaxlib/pjit.cc:668] ComputeCallSignature failed: (_reduce_prod) :INVALID_ARGUMENT: Not supported: The C++ ToPyArgSignature only accepts Buffer/DeviceArray, Numpy arrays scalars of supported types (see implementation), or Python scalars. Got type <class 'jax._src.interpreters.ad.JVPTracer'>
I jaxlib/pjit.cc:668] ComputeCallSignature failed: (_where) :INVALID_ARGUMENT: Not supported: The C++ ToPyArgSignature only accepts Buffer/DeviceArray, Numpy arrays scalars of supported types (see implementation), or Python scalars. Got type <class 'jax._src.interpreters.partial_eval.DynamicJaxprTracer'>

Notably, you get these errors on JAX's built-in functions as well - so even if you jax.jit(jax.vmap(f)) where f is just a series of built-in ops, you get all these failures from pjit.cc - the code still runs, and I assume runs correctly all on GPU with a nicely compiled/optimised kernel, but I've been trying to figure out why recent updates have caused some code I have to massively lose performance.

Which, looking at pjit.cc:

jax/jaxlib/pjit.cc

Lines 641 to 657 in 8ff5d33

	// Perform a few checks for the arguments. Currently we are only allowing
	// committed PyArray inputs. For other cases, e.g. Tracers or ShapedArray, it
	// will fallback to python. For jit, numpy arrays and scalars are also
	// allowed, which we will check later.
	for (const auto& arg : flat_dynamic_args) {
	if (arg.type().ptr() != xla::PyArray::type().ptr()) {
	continue;
	}
	xla::PyArray py_array = nb::borrow<xla::PyArray>(arg);
	// Only allow committed PyArray in cpp pjit for now as the logic on handling
	// sharding for uncommitted PyArray is complicated and still under
	// development.
	//
	// TODO(chky): Consider support uncommitted PyArray in cpp when the python
	// side stabilizes.

seems to be the correct behaviour - I assume the outer level compilation is still completing successfully though?

You can produce this behaviour with this:

import jax
import jax.numpy as jnp
import numpy as np
jax.config.update("jax_platform_name", "gpu") 

@jax.jit
def simple_stuff(x, y):
    z = jnp.matmul(x, y)
    z = jnp.abs(z)
    z = z**2
    return jnp.sum(z)

many_x = jnp.array(np.random.random((batch_dim, 1000, 1000)))
many_y = jnp.array(np.random.random((batch_dim, 1000, 1000)))

def do_many(x, y):
    res = simple_stuff(x,y) - 1
    return res

v_matmul = jax.jit(jax.vmap(do_many, in_axes=(0,0)))
res = v_matmul(many_x, many_y)

Which seems to indicate to me that true nested jit should be avoided - unless of course you use some of the nested functions independently elsewhere where they won't be passed Tracers as input. From https://docs.jax.dev/en/latest/faq.html#benchmarking-jax-code :

To get maximum performance from JAX, you should apply jax.jit() on your outer-most function calls.

If my understanding of all of this is correct, I'd love if someone with authority can confirm/correct!