what is going on when "recompilation" happens in a for loop

[google-deepmind]

what is going on when "recompilation" happens in a for loop

[cdoersch]

This is standard behavior of JAX. A function that's wrapped in jit will be compiled the first time it runs. Then the compiled version will be cached. Thus, if the jitted function is called with arguments of the same sizes and types, then it will use the cached version. Otherwise it needs to recompile. Every function in python has associated state (i.e., jax.jit returns a function which is a closure), so references to the compiled computation graph are stored in the state associated with the function.

In the future, this kind of question might be a better fit for a the JAX repository, as this behavior isn't specific to TAP/TAPIR.

[CHYjeremy]

Thanks @cdoersch , i will pay close attention to where i post issues to. Thank you for your explanations on it, clear and concise.