Slicing an array inside a JIT compiled function vs using sliced arrays as inputs to JIT compiled function

[JaySandesara]

Hello,

I have a function as follows that needs to update the values of ratio_x and ratio_y:

@jax.jit
def func(param_array, arr_ratio_up_x, arr_ratio_down_x, arr_ratio_up_y, arr_ratio_down_y):

    mu = param_array[0]
    mu_0 = param_array[1]
    mu_1 = param_array[2]
    mu_2 = param_array[3]

    def calculate_variations(n, carry):

        ratio_x, ratio_y = carry

        ratio_x *= jax.lax.cond(jnp.abs(param_array[n])<=1.0,poly_interp,exp_interp,(param_array[n], arr_ratio_up_x[n-4], arr_ratio_down_x[n-4], 1.0))

        ratio_y *= jax.lax.cond(jnp.abs(param_array[n])<=1.0,poly_interp,exp_interp,(param_array[n], arr_ratio_up_y[n-4], arr_ratio_down_y[n-4], 1.0))

        return (ratio_x, ratio_y)

    ratio_x, ratio_y = jax.lax.fori_loop(4, len(param_array), calculate_variations, tuple_init)

    sigma = sigma(mu, mu_0, mu_1, mu_2, ratio_x, ratio_y)

    return sigma

Here I have simplified my original function quite a bit, but the gist is that I have these repeated calls to poly_interp/exp_interp functions for the update of the different ratio_* variables. Note that arr_ratio_up_x and arr_ratio_up_y are not necessarily the same shape.

Since I have O(50) such ratio_* variables to update, I figured a faster way would be to just call the poly_interp/exp_interp function a single time in a single iteration. To this end, I create a combined input array using:

combined_var_up = jnp.column_stack((arr_ratio_up_x, arr_ratio_up_y))
combined_var_down = jnp.column_stack((arr_ratio_down_x, arr_ratio_down_y))
len_x = arr_ratio_up_x.shape[0]
len_y = arr_ratio_up_y.shape[0]

and then rewriting the function as follows:

@jax.jit
def func_optimized(param_array, combined_var_up, combined_var_down):
    
    
    mu = param_array[0]
    mu_qqZZ_0 = param_array[1]
    mu_qqZZ_1 = param_array[2]
    mu_qqZZ_2 = param_array[3]

    def calculate_variations(n, carry):

        combined_var_array_temp = carry

        combined_var_array_temp *= jax.lax.cond(jnp.abs(param_array[n])<=1.0, 
                                      poly_interp, 
                                      exp_interp, 
                                      (param_array[n], combined_var_up[n-4], combined_var_down[n-4], tuple_init_alt))   

        return combined_var_array_temp


    combined_var_array = jax.lax.fori_loop(4, len(param_array), calculate_variations, tuple_init_alt)
    
    ratio_x = jax.lax.slice(combined_var_array, (0,),(len_x,))
    ratio_y = jax.lax.slice(combined_var_array, (len_x,),(len_x+len_y,))
    
    sigma = sigma(mu, mu_0, mu_1, mu_2, ratio_x, ratio_y)

    return sigma

Since the updates happen in a single array-based computation, I figured this would be a more efficient jit-compiled function. But rather this is much slower in the compilation step and seems to crash the Jupyter kernel likely due to memory issues. I suspect the culprit is the slicing of the array combined_var_array? Is there a better way to slice them or use them in the function sigma? Or is using the original function func a better approach than func_optimized? Note that len_x and len_y are fixed at compile time and dont change.

Regards,
Jay Sandesara

[JaySandesara]

Okay, I found that the problem is only when I use a wrapper function to call the func_optimized function:

@jax.jit
def func_optimized_wrapper(tuple):

    return func_optimized(tuple, combined_var_up, combined_var_down)

If I call func_optimized directly, the compilation/execution times are fairly similar to func (is this expected?). I need to use the wrapper function func_optimized_wrapper which is then passed to another optimization program that requires this particular type of input. Any suggestions on what might be causing this issue with the wrapper function and not directly with the original function?

[JaySandesara]

Okay, found the solution, if I define the wrapper function with @jax.jit decorator then things dont work:

@jax.jit
def func_optimized_wrapper(tuple):

    return func_optimized(tuple, combined_var_up, combined_var_down)

But if I just remove @jax.jit in the wrapper, things work as expected:

def func_optimized_wrapper(tuple):

    return func_optimized(tuple, combined_var_up, combined_var_down)

I keep the @jax.jit in the original func_optimized function.

I am still curios for the explanation behind this, so I will keep the discussion open for now. It is possibly a new update, since I recently upgraded to v0.4.19 from v0.4.6. Also curios as to why the jit-compiled func_optimized with its single array computations is not faster than the jit-compiled func with its repeated function calls. Let me know if you have any insights/suggestions.

[jakevdp]

I'm not sure offhand what might be causing the slow execution. A minimal reproducible example would help in being able to better answer the question (it's much easier to debug behavior when you can see it yourself!).

One possible idea: rather than a fori_loop with repeated indexing, you might have better luck with jax.lax.scan, which directly scans over an axis of an array.

[JaySandesara]

Thanks @jakevdp for the prompt response! I have created a sample notebook (link) with simple inputs for reproducibility. Does this work for you?

[jakevdp]

Thanks - there's a lot there. I'm not entirely sure what your main question is. But many of the observations seem reasonable. for example this:

#Far Longer execution time when used with JIT-compilation
%timeit -r 1 -n 1 func_wrapper(param_array)

It seems reasonable that a function marked for jit-compilation would take longer to compile than a function not marked for jit-compilation, because a non-jit function does not incur any compilation cost.

[JaySandesara]

Hey @jakevdp, so yes you are correct that the compilation step would take longer. However, my comment was the time it takes to compile the wrapper vs the time it takes to compile the non-wrapped function has a ~2x difference. I would assume that since the wrapper simply calls the original function it should be the same? I have found that it also uses up much more memory if I jit compile the wrapper function, crashing the kernel for my original use case. Could this have to do with constant-folding?

And my question is why using array-based computation with func_optimized doesnt have too big an impact on compilation/execution times. And if there is a better solution to write this function more efficiently? This becomes more of a problem in my full use case where there are O(50) ratio_* variables to update in each loop. Then, the compilation time and post-compilation execution times reach the same order of magnitude. I tried lax.scan but it brings marginal improvements. If you have any other suggestion, it would be very helpful!

[jakevdp]

Oh I see – the difference is that when you compile func, you are passing all the arguments past the first one as arguments to the function, and so they are treated as runtime traced values. When you compile func_wrapper, you are closing over all the variables after the first one, and their global values are stored as numpy arrays, so they are being treated as compile-time constants. This changes the compilation characteristics: in particular the compiler can often specialize execution if a value is a constant, but sometimes this can lead to longer compilation times because the compiler has more flexibility, and more potential optimizations to explore. Does that answer your question?

[JaySandesara]

Yes, that is very helpful thanks! Do you have any suggestions on how I can optimize the function further? If not, we can close this discussion.

[jakevdp]

To be honest, most of this discussion strikes me as premature optimization. Fiddling with whether inputs are tuples or individual arguments, or whether arguments are passed to the function or closed over, will not likely lead to appreciable performance differences. I would focus on the algorithm itself rather than how you package it.