JIT with closure function

[Qiustander]

Hi all, I have a function generator which is like this

def func_generator(base_fn: Callable) -> Callable:
    def inner(parameters: jax.Array) -> Callable[jax.Array, jax.Array]:
          comp_fn  = base_fn(parameters)
          return comp_fn  
    return inner

That is, I have a base_fn, then it is relate to parameters. base_fn(parameters) will return another function that for usage. Then I compile_fn = jit(func_generator(base_fn)) and input parameters each time to get the functions for computation in other places.

My question is, will it lead to retracing each time I input a new parameters since comp_fn is, though inside the compilation, a function that JAX could not trace. How to check whether it is retracing?

[jakevdp]

Do you have a typo in your code? It appears the function is trying to return a reference to itself.

[Qiustander]

Oh, quite sorry about the typo, should return inner instead of func_generator. Fixed it

[jakevdp]

The answer to whether this will trigger a recompilation each time it is called depends on exactly when and where you're using jit. The key is that jit caches compilations based on the object id of the function object passed to jit. As written, func_generator will return a new callable (with a new id) each time you call it.

So with that in mind, if you do something like this, it will always re-trace and re-compile on the second call:

jax.jit(func_generator(base_fn))(params)
jax.jit(func_generator(base_fn))(params2)  # re-traces always

If you do something like this, you avoid creating another instance of the closed-over function, so the second call will hit the jit cache, and it will not lead to re-tracing when the inputs are compatible (i.e. params2 has the same shape and dtype as params).

jit_f = jax.jit(func_generator(base_fn))
jit_f(params)
jit_f(params2)  # does not re-trace if params2 is compatible with params

[Qiustander]

Thanks for the quick reply!

I am confused about one statement you write: As written, func_generator will return a new callable (with a new id) each time you call it. Yes, it will return a new id in each call, then why jit_f(params2) # does not re-trace if params2 is compatible with params? Since I create a new id inside the jitted function would not lead to re-compile?

My actual function is a bit more complex than it, is it possible to use some tools to check whether re-tracing happens or not?

The answer to whether this will trigger a recompilation each time it is called depends on exactly when and where you're using jit. The key is that jit caches compilations based on the object id of the function object passed to jit. As written, func_generator will return a new callable (with a new id) each time you call it.

So with that in mind, if you do something like this, it will always re-trace and re-compile on the second call:

jax.jit(func_generator(base_fn))(params)
jax.jit(func_generator(base_fn))(params2)  # re-traces always

If you do something like this, you avoid creating another instance of the closed-over function, so the second call will hit the jit cache, and it will not lead to re-tracing when the inputs are compatible (i.e. params2 has the same shape and dtype as params).

jit_f = jax.jit(func_generator(base_fn))
jit_f(params)
jit_f(params2)  # does not re-trace if params2 is compatible with params

[jakevdp]

When you call func_generator(base_fn), it returns a new function. jit_f wraps that function, and calling jit_f does not call func_generator, it calls the function previously returned by func_generator.

[Qiustander]

Thanks, got it! The inner is already cached right? So even it generates another function object comp_fn, which is inside the inner, it would not lead to re-compile as it caches all operations inside the jit_f.

A small follow-up question, I used jit outside a big for-loop, and each iteration I call func_generator to generate inner functions, and calculate something via putting parameters to jit_f. The base_fn is wrapper ouside the for-loop using closure function like the one shown above. I am not sure whether this design will lead to problem, so do you have suggestion to check whether retracing/memory problem happens since each time I generate new function objects inside?

Thanks, got it! The inner is already cached right? So even it generates another function object comp_fn, which is inside the inner, it would not lead to re-compile as it caches all operations inside the jit_f.

A small follow-up question, I used jit outside a big for-loop, and each iteration I call func_generator to generate inner functions, and calculate something via putting parameters to jit_f. The base_fn is wrapper ouside the for-loop using closure function like the one shown above. I am not sure whether this design will lead to problem, so do you have suggestion to check whether retracing/memory problem happens since each time I generate new function objects inside?

[jakevdp]

The inner is already cached right?

The inner function is the return value of func_generator, so when you JIT-compile that return value you are JIT-compiling (and cacheing) inner. Does that make sense?

I used jit outside a big for-loop, and each iteration I call func_generator

If you have everything under an outer JIT, then the relevant JIT cache / compilation is for that outer-jit function. Each inner function call will be inlined into the jaxpr of that outer function, and that will be compiled once. There will not be any cache for functions created within a JIT-compiled function.

In general there is no issue with this approach, but keep in mind that if you wrap a for loop in jit, that for loop will be unrolled during tracing before being passed to XLA. Depending on the size of the loop, this can lead to very long programs that might lead to slow compilation times.