Jitted function is fast to complete computations, but slow to return

[throwaway0123456]

Hi,

Firstly, apologies in advance for not providing a real code example. I have spent a long time trying to replicate my issue in a short snippet of code, but have so far failed to do so, any short example has not replicated this behaviour.

I will try to keep the question general, and hopefully the behaviour is explainable with the details I provide. Here is some simple pseudocode to illustrate the issue I am having:

@jax.jit
def jitted_func(A, b):
    x = (solve for x in Ax = b using lineax library)
    print('X has been calculated. Now function is about to return', x)
    return x

x = jitted_func(matrix, vector)
print('x has been calculated', x)

I am experiencing some strange behaviour where my jitted function runs very quickly right up until the return, but returning takes up to 2 minutes to complete. I have verified this by seeing that the first print happens instantly, but there is a delay of ~2 mins to see the second return, as in the pseudocode^.

Within this jitted function, I am using the lineax library which is a nice library for solving linear systems such as Ax = b. I thought that it is possible that this library is not jittable, however, what I find strange is that all operations within the function complete quickly, but it is returning from the function that takes long (~two minutes).

Finally, to add one detail which has particularly confused me, is that if I add one final operation to the end of my jitted function:
x = tree_map(lambda y: jnp.zeros_like(y), x), then my jitted function runs fine without getting stuck on the return. I have confirmed that the jax.make_jaxpr compiled code does not remove the previously applied lineax solve operations and replace them with this zeros_like operation. Therefore, adding an operation where I replace my FrozenDict with zeros somehow avoids freezing on returning from a function, when passing the real solution does freeze.

I have spent a lot of time trying to understand why this behaviour happens but I'm struggling to understand how to even begin to approach it. When trying to Jit the function that all of this code lives within, then I experience the 2 minute delay on returning from the function 1 level out. If I avoid jitting my function, then I don't experience this 2 minute lag on returning my function, but I instead have to deal with the solving of Ax=b being over 2 minutes.

I hope that my explanation was clear, please don't hesitate to ask for any other details. I understand this may be a very niche issue and I may have missed some detail that is crucial in understanding the behaviour. I myself have not yet even been able to replicate it within a single short snippet of code. Thank you for your time.

[jakevdp]

The reason for this is that the print statement executes during tracing, and it's not until tracing is finished that the function is compiled and then executed. See How to think in JAX for some general background on JAX's computational model.

[throwaway0123456]

Thank you @jakevdp. Apologies for the long post, that makes much more sense now.

One last point I am not clear on is why the following code:

@jax.jit
def jitted_func(A, b):
    x = (solve for x in Ax = b using lineax library)
    x = tree_map(lambda y: jnp.zeros_like(y), x)
    return x

x = jitted_func(matrix, vector)

runs suspiciously fast (2E-4 seconds), where x is replaced with zeros before returning. An obvious explanation is that jax is smart enough to understand that solving for x does not change the outcome of the function, and it should instead be filled with 0s right away. However, I checked the make_jaxpr for both versions of the function, and both involve the same component of solving for x. As I understand, both should take just as long to run.
Do you have any idea about why this code would run much faster?

Thanks again for your help.

[jakevdp]

It looks like the lineax computation doesn’t contribute to the function output, so the compiler eliminates it.

[throwaway0123456]

Thank you @jakevdp for your help on this issue, it has now been resolved.

I've now run into a potentially related issue but I'm not able to make sense of it with what you provided above. Let me know if I should make a separate thread, as I'm not sure if it is related or not. I have a jitted function derivative_fn. I find that measuring the time it takes to execute this in a for loop, is somehow much longer than if I nest it within another (non-jitted) function:

This following code set-up gives expected results:

def timing_function(func):
    start = time.time()
    inner_deriv, outer_deriv = func
    print('Func time ', time.time() - start)
    return inner_deriv, outer_deriv

...
...
...

starter = time.time()
for i in range(iterations):
    inner, outer = timing_function(derivative_fn(inner_params, outer_params, inner_kwargs, outer_kwargs))
    print('Loop time ', time.time() - starter)
    starter = time.time()

In the above code, I'm evaluating the derivative_fn inside the timing_function. When I print the "Func time", it reports only 1e-7, while printing the "Loop time" gives approx 1e-2. I'm wondering where this loss in time is coming from.

This appears to be a similar issue to what I posted above, where it seems like code is getting "stuck" on the return from a function. However, unlike in the previous case, my timing_function is not jitted, where the tracer pass would explain this apparent time difference.

Running the derivative_fn directly in the for loop again results in times of 1e-2, yet I have proof of the true jitted computation taking 1e-7 within timing_function - I have printed the outputs and verified that the calculation is actually being completed in 1e-7.

Thanks again for all of your help, it has been really helpful to understand better how Jax and jitting works.

[jakevdp]

You are evaluating derivative_fn before passing the result to timing_function. So within timing_function you are measuring how long it takes to unpack the returned tuple into two elements, and within "loop time" you are measuring the actual execution time of the function.

[throwaway0123456]

Apologies, that's my bad, I don't know how I didn't realise I was doing that. Thanks again for all of your help!