how to measure compile time without evaluation?

[phc27x]

I feel this must have been asked/answered a few times but I can't find anything, sorry if duplicate.

How can I compile a jitted-function without also evaluating it at the same time?
Eg, Is there a way I can pass the shape only of the input, to trigger the tracing & compiling, but without then triggering an execution.

At the moment, to measure compile time, I end up executing twice and looking at the difference between first and second execution, but this seems inefficient compared to being able to isolate the compilation step directly.

Thanks

[hawkinsp]

Try:

In [5]: import jax

In [6]: f = jax.jit(lambda x: x * 2)

In [7]: %timeit f.lower(42).compile()
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
7.54 ms ± 231 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

.lower() is new and not yet documented, but it does exactly what you want.

[phc27x]

This is exactly what I'm looking for, thanks!