Displaying progress of JAX lowering and compilation

[cool-RR]

Hi,

My program takes anywhere from 2 minutes to 10 hours to lower and compile my JAX code. It's frustrating that I don't know in advance how long I'll have to wait. Is there some way to measure the progress of the lowering and compilation? If so, I might be able to hook a tqdm progress bar to it and make the user experience much nicer.

Thanks,
Ram.

[ddrous]

Hi @cool-RR, 10 hours sounds really long : )! Perhaps some of the more savvy JAX folks know of a way to measure such timings.

In my experience, such variable compile times tend to be due to unrolled for loops in which the terminal value n varies across your runs. If that is indeed the case, then fori_loop could be a day-saver ?

Cheers

[jakevdp]

Hi @cool-RR – I don't know of any way to get a tqdm-style progress bar on compilation progress. If you're getting multi-hour compilations, I suspect your effort would be better spent changing your code to improve compilation speed. As @ddrous mentioned, the main culprit in cases like this is often Python control flow (e.g. for loops, while loops, or recursion) within a jax.jit-compiled function. Python control flow is flattened during tracing and passed to the compiler, and the compiler is typically super-linear in the number of operations in the program. So if you have e.g. a jit-compiled for-loop, you might replace it with either fori_loop or scan, or if each step of the loop is relatively costly, you could use a for loop over jit-compiled operations.

I hope that helps!

[cool-RR]

Thank you. Maybe I indeed have huge for-loops in there. The problem is I've got lots of functions calling other functions so I'm not sure which of these for-loops end up being huge. Does JAX provide any way for me to see which of my for-loops end up being huge when compiled?

[jakevdp]

Since for loop bounds must be static, that's usually pretty straightforward to see from the code: e.g. if you find for i in range(x.shape[0]), and x is large, then you know it's a large for-loop.

[cool-RR]

If I have a for-loop that only runs 2-3 times, but there's lots of code happening in the loop body, do I stand to gain a meaningful improvement by refactoring it into a fori_loop?